bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2FusedKernels.swift

@@ -0,0 +1,198 @@
+// Flux2FusedKernels.swift - Custom Metal kernels for performance-critical operations
+// Copyright 2025 Vincent Gourbin
+//
+// The RoPE kernel implementation is inspired by mzbac/flux2.swift
+// (https://github.com/mzbac/flux2.swift), licensed under Apache 2.0.
+// This implementation has been adapted for [B, H, S, D] tensor layout
+// (original uses [B, S, H, D]) and integrated with our caching system.
+
+import Foundation
+import MLX
+
+/// Collection of fused Metal kernels for performance-critical transformer operations
+///
+/// These kernels fuse multiple operations into a single GPU dispatch, reducing
+/// kernel launch overhead and improving memory locality.
+public enum Flux2FusedKernels {
+    /// Set to true to disable fused kernels (for debugging/comparison)
+    /// Can also be set via environment variable FLUX2_DISABLE_FUSED_KERNELS=1
+    nonisolated(unsafe) public static var disabled: Bool = {
+        ProcessInfo.processInfo.environment["FLUX2_DISABLE_FUSED_KERNELS"] == "1"
+    }()
+    /// Number of attention heads processed per thread for RoPE kernel
+    private static let ropeHeadsPerThread = 4
+
+    /// Custom Metal kernel for applying rotary position embeddings
+    ///
+    /// This kernel fuses the rotation operation (x * cos + rotated_x * sin) into a single
+    /// GPU kernel, processing multiple heads per thread for efficiency.
+    ///
+    /// Input layout: [B, H, S, D] where B=batch, H=heads, S=sequence, D=head_dim
+    /// Cos/Sin layout: [S, D] or [B, S, D]
+    private static let ropeKernel = MLXFast.metalKernel(
+        name: "flux2_apply_rope_bhsd",
+        inputNames: ["x", "cos", "sin", "batch", "heads", "seq", "head_dim", "rotary_batch_stride"],
+        outputNames: ["out"],
+        source: """
+            uint dp = thread_position_in_grid.x;           // dimension pair index (0..half_dim-1)
+            uint seq_idx = thread_position_in_grid.y;      // sequence position
+            uint z = thread_position_in_grid.z;            // batch * head_groups
+
+            uint half_dim = uint(head_dim / 2);
+            if (dp >= half_dim || seq_idx >= uint(seq)) {
+                return;
+            }
+
+            int head_groups = (heads + (N - 1)) / N;
+            int batch_idx = int(z) / head_groups;
+            int head_group = int(z) - batch_idx * head_groups;
+            if (batch_idx >= batch) {
+                return;
+            }
+
+            int head_base = head_group * N;
+
+            // cos/sin layout: [S, D] if rotary_batch_stride=0, else [B, S, D]
+            ulong row_base = (ulong(seq_idx) + ulong(batch_idx) * ulong(rotary_batch_stride)) * ulong(head_dim);
+            device const float* cos_row = cos + row_base;
+            device const float* sin_row = sin + row_base;
+
+            uint d0 = dp * 2;
+            float c = cos_row[d0];
+            float s = sin_row[d0];
+
+            // For [B, H, S, D] layout:
+            // x[batch][head][seq][dim] = x[batch * H * S * D + head * S * D + seq * D + dim]
+            ulong stride_S = ulong(head_dim);
+            ulong stride_H = ulong(seq) * stride_S;
+            ulong stride_B = ulong(heads) * stride_H;
+
+            for (int i = 0; i < N; ++i) {
+                int head_idx = head_base + i;
+                if (head_idx >= heads) {
+                    break;
+                }
+
+                ulong base = ulong(batch_idx) * stride_B + ulong(head_idx) * stride_H + ulong(seq_idx) * stride_S + ulong(d0);
+                float x0 = float(x[base]);
+                float x1 = float(x[base + 1]);
+
+                // Apply rotation: [x0, x1] * [[c, -s], [s, c]]
+                // Result: [x0*c - x1*s, x0*s + x1*c]
+                // But diffusers uses: x * cos + rotate_half(x) * sin
+                // where rotate_half gives [-x1, x0] for consecutive pairs
+                // So: [x0, x1] -> [x0*c - x1*s, x1*c + x0*s]
+                out[base] = T(x0 * c - x1 * s);
+                out[base + 1] = T(x1 * c + x0 * s);
+            }
+            """
+    )
+
+    /// Apply rotary embeddings using the fused Metal kernel
+    ///
+    /// - Parameters:
+    ///   - x: Input tensor [B, H, S, D] in bfloat16, float16, or float32
+    ///   - cos: Cosine embeddings [S, D] or [B, S, D] in float32
+    ///   - sin: Sine embeddings [S, D] or [B, S, D] in float32
+    /// - Returns: Rotated tensor [B, H, S, D] in same dtype as input, or nil if kernel can't be used
+    public static func applyRotaryEmb(
+        _ x: MLXArray,
+        cos: MLXArray,
+        sin: MLXArray
+    ) -> MLXArray? {
+        // Allow disabling for debugging
+        guard !disabled else { return nil }
+
+        // Only use kernel on GPU
+        guard Device.defaultDevice().deviceType == .gpu else {
+            return nil
+        }
+
+        // Validate dtypes
+        guard x.dtype == .bfloat16 || x.dtype == .float16 || x.dtype == .float32,
+              cos.dtype == .float32,
+              sin.dtype == .float32
+        else {
+            return nil
+        }
+
+        // Validate x shape: [B, H, S, D]
+        guard x.ndim == 4 else {
+            return nil
+        }
+
+        let batch = x.dim(0)
+        let heads = x.dim(1)
+        let seq = x.dim(2)
+        let headDim = x.dim(3)
+
+        // Validate cos/sin shape: [S, D] or [B, S, D]
+        guard (cos.ndim == 2 || cos.ndim == 3), cos.ndim == sin.ndim else { return nil }
+
+        let batchStride: Int32
+        if cos.ndim == 2 {
+            // [S, D] - shared across batch
+            guard cos.dim(0) == seq, sin.dim(0) == seq else { return nil }
+            guard cos.dim(1) == headDim, sin.dim(1) == headDim else { return nil }
+            batchStride = 0
+        } else if cos.ndim == 3 {
+            // [B, S, D] - per-batch embeddings
+            guard cos.dim(0) == batch, sin.dim(0) == batch else { return nil }
+            guard cos.dim(1) == seq, sin.dim(1) == seq else { return nil }
+            guard cos.dim(2) == headDim, sin.dim(2) == headDim else { return nil }
+            batchStride = Int32(seq)
+        } else {
+            return nil
+        }
+
+        // Head dimension must be even for rotation
+        guard headDim % 2 == 0 else {
+            return nil
+        }
+
+        let halfDim = headDim / 2
+        let headsPerThread = ropeHeadsPerThread
+        let headGroups = (heads + headsPerThread - 1) / headsPerThread
+
+        // Thread group sizing
+        let tgX = max(min(32, halfDim), 1)
+        let tgY = max(min(4, seq), 1)
+
+        let out = ropeKernel(
+            [
+                x,
+                cos,
+                sin,
+                Int32(batch),
+                Int32(heads),
+                Int32(seq),
+                Int32(headDim),
+                batchStride,
+            ],
+            template: [
+                ("T", x.dtype),
+                ("N", headsPerThread),
+            ],
+            grid: (halfDim, seq, batch * headGroups),
+            threadGroup: (tgX, tgY, 1),
+            outputShapes: [x.shape],
+            outputDTypes: [x.dtype],
+            stream: .gpu
+        )
+        return out[0]
+    }
+
+    /// Check if the fused RoPE kernel is available for the given configuration
+    public static func isRopeKernelAvailable(
+        xDtype: DType,
+        xNdim: Int,
+        cosNdim: Int,
+        sinNdim: Int
+    ) -> Bool {
+        guard Device.defaultDevice().deviceType == .gpu else { return false }
+        guard xDtype == .bfloat16 || xDtype == .float16 || xDtype == .float32 else { return false }
+        guard xNdim == 4 else { return false }
+        guard (cosNdim == 2 || cosNdim == 3), cosNdim == sinNdim else { return false }
+        return true
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2KVCache.swift

@@ -0,0 +1,36 @@
+// Flux2KVCache.swift - KV Cache for Klein 9B acceleration
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+
+/// KV Cache for a single attention layer
+/// Stores the key and value projections for reference tokens
+public struct Flux2KVCacheLayer {
+    public let k: MLXArray  // [B, numHeads, refSeqLen, headDim]
+    public let v: MLXArray  // [B, numHeads, refSeqLen, headDim]
+
+    public init(k: MLXArray, v: MLXArray) {
+        self.k = k
+        self.v = v
+    }
+}
+
+/// Complete KV Cache for all transformer layers
+/// Used to skip reference token re-processing after step 0
+public struct Flux2KVCache {
+    public let doubleStreamLayers: [Flux2KVCacheLayer]    // 8 layers
+    public let singleStreamLayers: [Flux2KVCacheLayer]    // 48 layers
+    public let layerCount: Int
+
+    public init(doubleStreamLayers: [Flux2KVCacheLayer], singleStreamLayers: [Flux2KVCacheLayer]) {
+        self.doubleStreamLayers = doubleStreamLayers
+        self.singleStreamLayers = singleStreamLayers
+        self.layerCount = doubleStreamLayers.count + singleStreamLayers.count
+    }
+
+    public var arrays: [MLXArray] {
+        doubleStreamLayers.flatMap { [$0.k, $0.v] } +
+            singleStreamLayers.flatMap { [$0.k, $0.v] }
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2Modulation.swift

@@ -0,0 +1,193 @@
+// Flux2Modulation.swift - AdaLN Modulation for Flux.2
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Modulation parameters for a single AdaLN operation
+public struct ModulationParams {
+    /// Shift parameter for centering
+    public let shift: MLXArray
+    /// Scale parameter for variance
+    public let scale: MLXArray
+    /// Gate parameter for residual connection
+    public let gate: MLXArray
+
+    public init(shift: MLXArray, scale: MLXArray, gate: MLXArray) {
+        self.shift = shift
+        self.scale = scale
+        self.gate = gate
+    }
+}
+
+/// Adaptive Layer Normalization Modulation
+///
+/// Generates (shift, scale, gate) parameters from timestep embeddings
+/// for adaptive layer normalization in transformer blocks.
+public class Flux2Modulation: Module, @unchecked Sendable {
+    let dim: Int
+    let numSets: Int
+    @ModuleInfo var linear: Linear
+
+    /// Initialize modulation layer
+    /// - Parameters:
+    ///   - dim: Model dimension (6144 for Flux.2)
+    ///   - numSets: Number of modulation parameter sets (e.g., 6 for double-stream, 3 for single)
+    public init(dim: Int, numSets: Int) {
+        self.dim = dim
+        self.numSets = numSets
+        // Each set has 3 parameters: shift, scale, gate
+        // So output dim = dim * 3 * numSets
+        // IMPORTANT: bias=false to match diffusers checkpoint
+        self.linear = Linear(dim, dim * 3 * numSets, bias: false)
+    }
+
+    /// Generate modulation parameters
+    /// - Parameter embedding: Timestep/conditioning embedding [B, dim]
+    /// - Returns: Array of ModulationParams, one for each set
+    public func callAsFunction(_ embedding: MLXArray) -> [ModulationParams] {
+        Flux2Debug.verbose("Modulation input shape: \(embedding.shape)")
+
+        // Apply SiLU activation before projection (common in DiT-style models)
+        let activated = silu(embedding)
+        Flux2Debug.verbose("After silu: \(activated.shape)")
+
+        // Project to all modulation parameters
+        let allParams = linear(activated)  // [B, dim * 3 * numSets]
+        Flux2Debug.verbose("After linear (dim=\(dim), numSets=\(numSets)): \(allParams.shape)")
+
+        let batchSize = allParams.shape[0]
+        let reshapedParams = allParams.reshaped([batchSize, numSets, 3, dim])
+        Flux2Debug.verbose("After reshape: \(reshapedParams.shape)")
+
+        // Split into individual parameter sets
+        var result: [ModulationParams] = []
+
+        for i in 0..<numSets {
+            let shift = reshapedParams[0..., i, 0, 0...]
+            let scale = reshapedParams[0..., i, 1, 0...]
+            let gate = reshapedParams[0..., i, 2, 0...]
+
+            Flux2Debug.verbose("Set \(i) - shift: \(shift.shape), scale: \(scale.shape), gate: \(gate.shape)")
+
+            result.append(ModulationParams(shift: shift, scale: scale, gate: gate))
+        }
+
+        return result
+    }
+
+    /// Generate modulation parameters as a flat array for efficient processing
+    /// - Parameter embedding: Timestep/conditioning embedding [B, dim]
+    /// - Returns: Tensor [B, numSets, 3, dim] containing all shift/scale/gate
+    public func forwardFlat(_ embedding: MLXArray) -> MLXArray {
+        let activated = silu(embedding)
+        let allParams = linear(activated)  // [B, dim * 3 * numSets]
+
+        // Reshape to [B, numSets, 3, dim]
+        let batchSize = allParams.shape[0]
+        return allParams.reshaped([batchSize, numSets, 3, dim])
+    }
+}
+
+/// Apply AdaLN modulation to normalized input
+/// - Parameters:
+///   - x: Normalized input [B, S, dim]
+///   - shift: Shift parameter [B, dim]
+///   - scale: Scale parameter [B, dim]
+/// - Returns: Modulated output [B, S, dim]
+public func applyModulation(
+    _ x: MLXArray,
+    shift: MLXArray,
+    scale: MLXArray
+) -> MLXArray {
+    Flux2Debug.verbose("applyModulation - x: \(x.shape), shift: \(shift.shape), scale: \(scale.shape)")
+
+    // Expand shift/scale for sequence dimension
+    // x: [B, S, dim], shift/scale: [B, dim]
+    let shiftExpanded = shift.expandedDimensions(axis: 1)  // [B, 1, dim]
+    let scaleExpanded = scale.expandedDimensions(axis: 1)  // [B, 1, dim]
+
+    Flux2Debug.verbose("Expanded - shift: \(shiftExpanded.shape), scale: \(scaleExpanded.shape)")
+
+    // Apply: x * (1 + scale) + shift
+    return x * (1 + scaleExpanded) + shiftExpanded
+}
+
+/// Apply gating to residual connection
+/// - Parameters:
+///   - residual: Residual to be gated [B, S, dim]
+///   - gate: Gate parameter [B, dim]
+/// - Returns: Gated residual [B, S, dim]
+public func applyGate(_ residual: MLXArray, gate: MLXArray) -> MLXArray {
+    let gateExpanded = gate.expandedDimensions(axis: 1)  // [B, 1, dim]
+    return residual * gateExpanded
+}
+
+/// Adaptive Layer Norm with continuous embedding (for final output)
+public class AdaLayerNormContinuous: Module, @unchecked Sendable {
+    let dim: Int
+    let norm: LayerNorm
+    @ModuleInfo var linear: Linear
+
+    public init(dim: Int, eps: Float = 1e-6, usesBias: Bool = false) {
+        self.dim = dim
+        self.norm = LayerNorm(dimensions: dim, eps: eps, affine: false)
+        self._linear = ModuleInfo(wrappedValue: Linear(dim, dim * 2, bias: usesBias))
+    }
+
+    /// Apply adaptive layer norm
+    /// - Parameters:
+    ///   - x: Input tensor [B, S, dim]
+    ///   - conditioning: Conditioning embedding [B, dim]
+    /// - Returns: Normalized and modulated output [B, S, dim]
+    public func callAsFunction(_ x: MLXArray, conditioning: MLXArray) -> MLXArray {
+        // Get scale and shift from conditioning
+        // IMPORTANT: diffusers order is linear(silu(x)), NOT silu(linear(x))
+        // And the split order is (scale, shift), NOT (shift, scale)
+        let params = linear(silu(conditioning))  // [B, dim * 2]
+        let scale = params[0..., 0..<dim]
+        let shift = params[0..., dim...]
+
+        // Normalize
+        let normalized = norm(x)
+
+        // Apply modulation
+        return applyModulation(normalized, shift: shift, scale: scale)
+    }
+}
+
+/// Modulation for double-stream blocks
+/// Generates 6 sets of parameters: 2 for image attention, 2 for image FFN, 2 for text
+public class DoubleStreamModulation: Module, @unchecked Sendable {
+    let imgModulation: Flux2Modulation
+    let txtModulation: Flux2Modulation
+
+    public init(dim: Int) {
+        // Image branch: attn (shift, scale, gate) + ffn (shift, scale, gate) = 2 sets * 3
+        self.imgModulation = Flux2Modulation(dim: dim, numSets: 2)
+        // Text branch: attn (shift, scale, gate) + ffn (shift, scale, gate) = 2 sets * 3
+        self.txtModulation = Flux2Modulation(dim: dim, numSets: 2)
+    }
+
+    public func callAsFunction(_ embedding: MLXArray) -> (img: [ModulationParams], txt: [ModulationParams]) {
+        let imgParams = imgModulation(embedding)
+        let txtParams = txtModulation(embedding)
+        return (img: imgParams, txt: txtParams)
+    }
+}
+
+/// Modulation for single-stream blocks
+/// Generates 2 sets: attention and FFN (fused in parallel attention)
+public class SingleStreamModulation: Module, @unchecked Sendable {
+    let modulation: Flux2Modulation
+
+    public init(dim: Int) {
+        // Single stream: combined attn+ffn modulation = 1 set * 3
+        self.modulation = Flux2Modulation(dim: dim, numSets: 1)
+    }
+
+    public func callAsFunction(_ embedding: MLXArray) -> [ModulationParams] {
+        modulation(embedding)
+    }
+}