bleam 0.0.8 → 0.0.10

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

@@ -0,0 +1,389 @@
+// Flux2ParallelAttention.swift - Fused Attention+FFN for Single-Stream Blocks
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Parallel Self-Attention for Single-Stream Blocks
+///
+/// This module fuses the QKV projection with MLP gate/up projections
+/// for more efficient computation. Used in the 48 single-stream blocks.
+///
+/// Architecture:
+/// - Single projection: dim -> (3*innerDim + 2*mlpHidden) for Q,K,V,gate,up
+/// - Self-attention on concatenated text+image
+/// - Output projection: (innerDim + mlpHidden) -> dim
+public class Flux2ParallelSelfAttention: Module, @unchecked Sendable {
+    let dim: Int
+    let numHeads: Int
+    let headDim: Int
+    let innerDim: Int
+    let mlpHiddenDim: Int
+
+    // Fused input projection: Q, K, V, MLP gate, MLP up
+    // @ModuleInfo allows update(modules:) to replace with LoRA
+    @ModuleInfo var toQkvMlp: Linear
+
+    // QK normalization
+    let normQ: FluxRMSNorm
+    let normK: FluxRMSNorm
+
+    // Fused output projection: attention output + MLP output -> dim
+    @ModuleInfo var toOut: Linear
+
+    /// Initialize Parallel Self-Attention
+    /// - Parameters:
+    ///   - dim: Model dimension (6144)
+    ///   - numHeads: Number of attention heads (48)
+    ///   - headDim: Dimension per head (128)
+    ///   - mlpRatio: MLP hidden dimension ratio (default 4.0)
+    public init(
+        dim: Int,
+        numHeads: Int,
+        headDim: Int,
+        mlpRatio: Float = 3.0
+    ) {
+        self.dim = dim
+        self.numHeads = numHeads
+        self.headDim = headDim
+        self.innerDim = numHeads * headDim
+        self.mlpHiddenDim = Int(Float(dim) * mlpRatio)
+
+        // Fused projection dimensions:
+        // Q: innerDim, K: innerDim, V: innerDim, gate: mlpHidden, up: mlpHidden
+        let totalProjDim = innerDim * 3 + mlpHiddenDim * 2
+        self._toQkvMlp.wrappedValue = Linear(dim, totalProjDim, bias: false)
+
+        // QK normalization
+        self.normQ = FluxRMSNorm(dim: headDim)
+        self.normK = FluxRMSNorm(dim: headDim)
+
+        // Output projection: attention (innerDim) + mlp (mlpHidden) -> dim (no bias)
+        self._toOut.wrappedValue = Linear(innerDim + mlpHiddenDim, dim, bias: false)
+    }
+
+    /// Forward pass
+    /// - Parameters:
+    ///   - hiddenStates: Combined image+text hidden states [B, S, dim]
+    ///   - rotaryEmb: Optional RoPE embeddings
+    /// - Returns: Updated hidden states [B, S, dim]
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil
+    ) -> MLXArray {
+        let batchSize = hiddenStates.shape[0]
+        let seqLen = hiddenStates.shape[1]
+
+        // Fused projection
+        let projected = toQkvMlp(hiddenStates)
+
+        // Split into Q, K, V, gate, up
+        var q = projected[0..., 0..., 0..<innerDim]
+        var k = projected[0..., 0..., innerDim..<(innerDim * 2)]
+        let v = projected[0..., 0..., (innerDim * 2)..<(innerDim * 3)]
+        let mlpGate = projected[0..., 0..., (innerDim * 3)..<(innerDim * 3 + mlpHiddenDim)]
+        let mlpUp = projected[0..., 0..., (innerDim * 3 + mlpHiddenDim)...]
+
+        // Reshape for multi-head attention: [B, S, H*D] -> [B, H, S, D]
+        q = reshapeForAttention(q, batchSize: batchSize, seqLen: seqLen)
+        k = reshapeForAttention(k, batchSize: batchSize, seqLen: seqLen)
+        let vReshaped = reshapeForAttention(v, batchSize: batchSize, seqLen: seqLen)
+
+        // Apply QK normalization
+        q = normQ(q)
+        k = normK(k)
+
+        // Apply RoPE if provided
+        if let rope = rotaryEmb {
+            (q, k) = applyRoPE(q: q, k: k, cos: rope.cos, sin: rope.sin)
+        }
+
+        // Compute attention
+        let attnOutput = MLXFast.scaledDotProductAttention(
+            queries: q,
+            keys: k,
+            values: vReshaped,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+
+        // Reshape attention output: [B, H, S, D] -> [B, S, H*D]
+        let attnOut = reshapeFromAttention(attnOutput, batchSize: batchSize, seqLen: seqLen)
+
+        // Compute MLP output: SwiGLU(gate, up)
+        let mlpOut = silu(mlpGate) * mlpUp
+
+        // Concatenate attention and MLP outputs
+        let combined = concatenated([attnOut, mlpOut], axis: -1)
+
+        // Final projection
+        return toOut(combined)
+    }
+
+    // MARK: - Helper Functions
+
+    /// Forward pass with cached reference K/V (for KV cache acceleration)
+    /// - Parameters:
+    ///   - hiddenStates: Image hidden states only [B, S_img, dim]
+    ///   - cachedRefKV: Pre-computed reference (text+ref image) K and V [B, H, S_ref, D]
+    ///   - rotaryEmb: Optional RoPE embeddings - only image portion needed
+    /// - Returns: Updated hidden states [B, S_img, dim]
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        cachedRefKV: (k: MLXArray, v: MLXArray),
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil
+    ) -> MLXArray {
+        let batchSize = hiddenStates.shape[0]
+        let seqLenImg = hiddenStates.shape[1]
+        let seqLenRef = cachedRefKV.k.shape[2]
+
+        // Fused projection for image tokens only
+        let projected = toQkvMlp(hiddenStates)
+
+        // Split into Q, K, V, gate, up
+        var q = projected[0..., 0..., 0..<innerDim]
+        var k = projected[0..., 0..., innerDim..<(innerDim * 2)]
+        let v = projected[0..., 0..., (innerDim * 2)..<(innerDim * 3)]
+        let mlpGate = projected[0..., 0..., (innerDim * 3)..<(innerDim * 3 + mlpHiddenDim)]
+        let mlpUp = projected[0..., 0..., (innerDim * 3 + mlpHiddenDim)...]
+
+        // Reshape for multi-head attention
+        q = reshapeForAttention(q, batchSize: batchSize, seqLen: seqLenImg)
+        k = reshapeForAttention(k, batchSize: batchSize, seqLen: seqLenImg)
+        let vReshaped = reshapeForAttention(v, batchSize: batchSize, seqLen: seqLenImg)
+
+        // Apply QK normalization
+        q = normQ(q)
+        k = normK(k)
+
+        // Apply RoPE to image Q/K only
+        if let rope = rotaryEmb {
+            let imgCos = rope.cos[seqLenRef..<(seqLenRef + seqLenImg)]
+            let imgSin = rope.sin[seqLenRef..<(seqLenRef + seqLenImg)]
+            (q, k) = applyRoPE(q: q, k: k, cos: imgCos, sin: imgSin)
+        }
+
+        // Concatenate cached reference K/V with image K/V
+        let concatK = concatenated([cachedRefKV.k, k], axis: 2)
+        let concatV = concatenated([cachedRefKV.v, vReshaped], axis: 2)
+
+        // Compute attention
+        let attnOutput = MLXFast.scaledDotProductAttention(
+            queries: q,
+            keys: concatK,
+            values: concatV,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+
+        // Reshape attention output
+        let attnOut = reshapeFromAttention(attnOutput, batchSize: batchSize, seqLen: seqLenImg)
+
+        // Compute MLP output
+        let mlpOut = silu(mlpGate) * mlpUp
+
+        // Concatenate and project
+        let combined = concatenated([attnOut, mlpOut], axis: -1)
+        return toOut(combined)
+    }
+
+    /// Compute K and V for reference tokens only (for KV cache building)
+    /// - Parameters:
+    ///   - hiddenStates: Normalized reference hidden states [B, S_ref, dim]
+    ///   - batchSize: Batch size
+    /// - Returns: (k, v) each [B, H, S_ref, D]
+    public func computeRefKV(
+        hiddenStates: MLXArray,
+        batchSize: Int
+    ) -> (k: MLXArray, v: MLXArray) {
+        let seqLenRef = hiddenStates.shape[1]
+
+        let projected = toQkvMlp(hiddenStates)
+        var k = projected[0..., 0..., innerDim..<(innerDim * 2)]
+        let v = projected[0..., 0..., (innerDim * 2)..<(innerDim * 3)]
+
+        k = reshapeForAttention(k, batchSize: batchSize, seqLen: seqLenRef)
+        let vReshaped = reshapeForAttention(v, batchSize: batchSize, seqLen: seqLenRef)
+
+        k = normK(k)
+
+        return (k: k, v: vReshaped)
+    }
+
+    // MARK: - Helper Functions
+
+    private func reshapeForAttention(_ x: MLXArray, batchSize: Int, seqLen: Int) -> MLXArray {
+        x.reshaped([batchSize, seqLen, numHeads, headDim])
+            .transposed(0, 2, 1, 3)
+    }
+
+    private func reshapeFromAttention(_ x: MLXArray, batchSize: Int, seqLen: Int) -> MLXArray {
+        x.transposed(0, 2, 1, 3)
+            .reshaped([batchSize, seqLen, numHeads * headDim])
+    }
+
+    private func applyRoPE(
+        q: MLXArray,
+        k: MLXArray,
+        cos: MLXArray,
+        sin: MLXArray
+    ) -> (MLXArray, MLXArray) {
+        let cosExpanded = cos.expandedDimensions(axes: [0, 1])
+        let sinExpanded = sin.expandedDimensions(axes: [0, 1])
+
+        let qRotated = rotateHalf(q)
+        let kRotated = rotateHalf(k)
+
+        let qOut = q * cosExpanded + qRotated * sinExpanded
+        let kOut = k * cosExpanded + kRotated * sinExpanded
+
+        return (qOut, kOut)
+    }
+
+    private func rotateHalf(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, S, D]
+        // Diffusers approach: reshape to [B, H, S, D/2, 2], then [-imag, real]
+        let batchSize = x.shape[0]
+        let numHeads = x.shape[1]
+        let seqLen = x.shape[2]
+        let dim = x.shape[3]
+        let halfDim = dim / 2
+
+        // Reshape to [B, H, S, D/2, 2]
+        let xReshaped = x.reshaped([batchSize, numHeads, seqLen, halfDim, 2])
+
+        // Get real and imag parts (consecutive pairs)
+        let xReal = xReshaped[0..., 0..., 0..., 0..., 0]  // [B, H, S, D/2]
+        let xImag = xReshaped[0..., 0..., 0..., 0..., 1]  // [B, H, S, D/2]
+
+        // Create rotated: stack [-imag, real] and flatten
+        let xRotatedStacked = stacked([-xImag, xReal], axis: -1)  // [B, H, S, D/2, 2]
+        return xRotatedStacked.reshaped([batchSize, numHeads, seqLen, dim])  // [B, H, S, D]
+    }
+}
+
+/// Alternative implementation with separate projections
+/// (for compatibility with some weight formats)
+public class Flux2ParallelSelfAttentionSplit: Module, @unchecked Sendable {
+    let dim: Int
+    let numHeads: Int
+    let headDim: Int
+    let innerDim: Int
+    let mlpHiddenDim: Int
+
+    // Separate projections (var for LoRA injection)
+    @ModuleInfo var toQ: Linear
+    @ModuleInfo var toK: Linear
+    @ModuleInfo var toV: Linear
+    @ModuleInfo var mlpGate: Linear
+    @ModuleInfo var mlpUp: Linear
+
+    // QK normalization
+    let normQ: FluxRMSNorm
+    let normK: FluxRMSNorm
+
+    // Separate output projections (var for LoRA injection)
+    @ModuleInfo var toAttnOut: Linear
+    @ModuleInfo var mlpDown: Linear
+
+    public init(
+        dim: Int,
+        numHeads: Int,
+        headDim: Int,
+        mlpRatio: Float = 3.0
+    ) {
+        self.dim = dim
+        self.numHeads = numHeads
+        self.headDim = headDim
+        self.innerDim = numHeads * headDim
+        self.mlpHiddenDim = Int(Float(dim) * mlpRatio)
+
+        self._toQ = ModuleInfo(wrappedValue: Linear(dim, innerDim, bias: false))
+        self._toK = ModuleInfo(wrappedValue: Linear(dim, innerDim, bias: false))
+        self._toV = ModuleInfo(wrappedValue: Linear(dim, innerDim, bias: false))
+        self._mlpGate = ModuleInfo(wrappedValue: Linear(dim, mlpHiddenDim, bias: false))
+        self._mlpUp = ModuleInfo(wrappedValue: Linear(dim, mlpHiddenDim, bias: false))
+
+        self.normQ = FluxRMSNorm(dim: headDim)
+        self.normK = FluxRMSNorm(dim: headDim)
+
+        self._toAttnOut = ModuleInfo(wrappedValue: Linear(innerDim, dim, bias: false))
+        self._mlpDown = ModuleInfo(wrappedValue: Linear(mlpHiddenDim, dim, bias: false))
+    }
+
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil
+    ) -> MLXArray {
+        let batchSize = hiddenStates.shape[0]
+        let seqLen = hiddenStates.shape[1]
+
+        // Separate projections
+        var q = toQ(hiddenStates)
+        var k = toK(hiddenStates)
+        let v = toV(hiddenStates)
+        let gate = mlpGate(hiddenStates)
+        let up = mlpUp(hiddenStates)
+
+        // Reshape for attention
+        q = q.reshaped([batchSize, seqLen, numHeads, headDim]).transposed(0, 2, 1, 3)
+        k = k.reshaped([batchSize, seqLen, numHeads, headDim]).transposed(0, 2, 1, 3)
+        let vReshaped = v.reshaped([batchSize, seqLen, numHeads, headDim]).transposed(0, 2, 1, 3)
+
+        // QK norm
+        q = normQ(q)
+        k = normK(k)
+
+        // RoPE
+        if let rope = rotaryEmb {
+            let cosExpanded = rope.cos.expandedDimensions(axes: [0, 1])
+            let sinExpanded = rope.sin.expandedDimensions(axes: [0, 1])
+
+            let qRot = rotateHalf(q)
+            let kRot = rotateHalf(k)
+
+            q = q * cosExpanded + qRot * sinExpanded
+            k = k * cosExpanded + kRot * sinExpanded
+        }
+
+        // Attention
+        let attnOut = MLXFast.scaledDotProductAttention(
+            queries: q,
+            keys: k,
+            values: vReshaped,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+
+        // Reshape back
+        let attnReshaped = attnOut.transposed(0, 2, 1, 3).reshaped([batchSize, seqLen, innerDim])
+
+        // MLP
+        let mlpOut = silu(gate) * up
+
+        // Separate output projections and add
+        return toAttnOut(attnReshaped) + mlpDown(mlpOut)
+    }
+
+    private func rotateHalf(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, S, D]
+        // Diffusers approach: reshape to [B, H, S, D/2, 2], then [-imag, real]
+        let batchSize = x.shape[0]
+        let numHeads = x.shape[1]
+        let seqLen = x.shape[2]
+        let dim = x.shape[3]
+        let halfDim = dim / 2
+
+        // Reshape to [B, H, S, D/2, 2]
+        let xReshaped = x.reshaped([batchSize, numHeads, seqLen, halfDim, 2])
+
+        // Get real and imag parts (consecutive pairs)
+        let xReal = xReshaped[0..., 0..., 0..., 0..., 0]  // [B, H, S, D/2]
+        let xImag = xReshaped[0..., 0..., 0..., 0..., 1]  // [B, H, S, D/2]
+
+        // Create rotated: stack [-imag, real] and flatten
+        let xRotatedStacked = stacked([-xImag, xReal], axis: -1)  // [B, H, S, D/2, 2]
+        return xRotatedStacked.reshaped([batchSize, numHeads, seqLen, dim])  // [B, H, S, D]
+    }
+}

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

@@ -0,0 +1,304 @@
+// Flux2RoPE.swift - Rotary Position Embeddings for Flux.2
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Cached RoPE embeddings with LRU eviction
+/// Note: Not Sendable because MLXArray isn't, but thread-safety is ensured by NSLock
+private struct RoPECacheEntry {
+    let cos: MLXArray
+    let sin: MLXArray
+    var lastAccess: UInt64
+}
+
+/// 4D Rotary Position Embedding for Flux.2
+///
+/// Flux.2 uses a 4D RoPE with axes_dims = [32, 32, 32, 32] for a total of 128 dims.
+/// This encodes position information for (T, H, W, L) axes.
+///
+/// Includes LRU caching to avoid recomputing embeddings for repeated position IDs.
+public class Flux2RoPE: Module, @unchecked Sendable {
+    let axesDims: [Int]
+    let theta: Float
+    let totalDims: Int
+
+    /// Maximum number of cached entries (LRU eviction when exceeded)
+    private let cacheMaxSize: Int
+
+    /// Cache for computed embeddings, keyed by shape hash
+    private var cache: [String: RoPECacheEntry] = [:]
+    private var accessCounter: UInt64 = 0
+    private let lock = NSLock()
+
+    /// Initialize RoPE with given axes dimensions
+    /// - Parameters:
+    ///   - axesDims: Dimensions per axis [T, H, W, L], default [32, 32, 32, 32]
+    ///   - theta: Base frequency for positional encoding, default 2000.0
+    ///   - cacheMaxSize: Maximum cache entries (default 8)
+    public init(
+        axesDims: [Int] = [32, 32, 32, 32],
+        theta: Float = 2000.0,
+        cacheMaxSize: Int = 8
+    ) {
+        self.axesDims = axesDims
+        self.theta = theta
+        self.totalDims = axesDims.reduce(0, +)
+        self.cacheMaxSize = cacheMaxSize
+    }
+
+    /// Compute frequency bands for given dimension
+    private func computeFreqs(dim: Int, maxLen: Int) -> MLXArray {
+        let halfDim = dim / 2
+        let freqSeq = MLXArray(0..<halfDim).asType(.float32)
+        let invFreq = 1.0 / pow(MLXArray(theta), freqSeq / Float(halfDim))
+
+        let positions = MLXArray(0..<maxLen).asType(.float32).expandedDimensions(axis: 1)
+        let freqs = positions * invFreq.expandedDimensions(axis: 0)
+
+        return freqs
+    }
+
+    /// Generate cache key from position IDs shape
+    /// Position IDs with the same shape can encode different image/reference layouts.
+    private func cacheKey(for ids: MLXArray) -> String {
+        var hash: UInt64 = 1469598103934665603
+        for value in ids.asType(.int32).asArray(Int32.self) {
+            hash ^= UInt64(bitPattern: Int64(value))
+            hash &*= 1099511628211
+        }
+
+        return "rope_shape=\(ids.shape)_axes=\(axesDims)_theta=\(theta)_hash=\(hash)"
+    }
+
+    /// Generate rotary embeddings for position IDs
+    /// - Parameter ids: Position IDs [S, 4] where 4 corresponds to (T, H, W, L)
+    /// - Returns: Tuple of (cos, sin) embeddings each of shape [S, totalDims]
+    public func callAsFunction(_ ids: MLXArray) -> (cos: MLXArray, sin: MLXArray) {
+        let key = cacheKey(for: ids)
+
+        // Check cache with lock
+        lock.lock()
+        if var entry = cache[key] {
+            accessCounter += 1
+            entry.lastAccess = accessCounter
+            cache[key] = entry
+            lock.unlock()
+            Flux2Debug.verbose("RoPE cache hit for \(key)")
+            return (cos: entry.cos, sin: entry.sin)
+        }
+        lock.unlock()
+
+        Flux2Debug.verbose("RoPE cache miss for \(key), computing...")
+
+        // Compute embeddings
+        let (cosEmb, sinEmb) = computeEmbeddings(for: ids)
+
+        // Store in cache with LRU eviction
+        lock.lock()
+        accessCounter += 1
+
+        // Evict LRU entry if cache is full
+        if cache.count >= cacheMaxSize {
+            if let lruKey = cache.min(by: { $0.value.lastAccess < $1.value.lastAccess })?.key {
+                cache.removeValue(forKey: lruKey)
+                Flux2Debug.verbose("RoPE cache evicted \(lruKey)")
+            }
+        }
+
+        cache[key] = RoPECacheEntry(cos: cosEmb, sin: sinEmb, lastAccess: accessCounter)
+        lock.unlock()
+
+        return (cos: cosEmb, sin: sinEmb)
+    }
+
+    /// Clear the RoPE embedding cache
+    public func clearCache() {
+        lock.lock()
+        cache.removeAll()
+        accessCounter = 0
+        lock.unlock()
+        Flux2Debug.verbose("RoPE cache cleared")
+    }
+
+    /// Compute embeddings (internal method)
+    private func computeEmbeddings(for ids: MLXArray) -> (cos: MLXArray, sin: MLXArray) {
+        var cosComponents: [MLXArray] = []
+        var sinComponents: [MLXArray] = []
+
+        let seqLen = ids.shape[0]
+
+        // Process each axis - INTERLEAVE EACH AXIS FIRST, THEN CONCATENATE
+        // This matches diffusers: repeat_interleave(2) per axis, then cat all axes
+        for (axisIdx, axisDim) in axesDims.enumerated() {
+            // Get position indices for this axis
+            let axisPositions = ids[0..., axisIdx]  // [S]
+
+            // Compute frequencies for this axis dimension
+            // diffusers: freqs = 1.0 / (theta ** (arange(0, dim, 2) / dim))
+            // For dim=32: arange(0, 32, 2) = [0, 2, 4, ..., 30] and /32 gives [0, 1/16, 2/16, ...]
+            let freqSeq = MLXArray(stride(from: 0, to: axisDim, by: 2)).asType(.float32)
+            let invFreq = 1.0 / pow(MLXArray(theta), freqSeq / Float(axisDim))
+
+            // Compute position * frequency
+            let posExpanded = axisPositions.asType(.float32).expandedDimensions(axis: 1)  // [S, 1]
+            let freqExpanded = invFreq.expandedDimensions(axis: 0)  // [1, halfDim]
+            let freqs = posExpanded * freqExpanded  // [S, halfDim]
+
+            // Compute cos and sin
+            let axisCosHalf = cos(freqs)  // [S, halfDim]
+            let axisSinHalf = sin(freqs)  // [S, halfDim]
+
+            // Interleave THIS axis: [c0, c0, c1, c1, ...] like repeat_interleave(2, dim=1)
+            // [S, halfDim] -> [S, halfDim, 1] -> [S, halfDim, 2] -> [S, axisDim]
+            let cosExpanded = axisCosHalf.expandedDimensions(axis: -1)  // [S, halfDim, 1]
+            let cosInterleaved = concatenated([cosExpanded, cosExpanded], axis: -1)  // [S, halfDim, 2]
+            let axisCos = cosInterleaved.reshaped([seqLen, axisDim])  // [S, axisDim]
+
+            let sinExpanded = axisSinHalf.expandedDimensions(axis: -1)  // [S, halfDim, 1]
+            let sinInterleaved = concatenated([sinExpanded, sinExpanded], axis: -1)  // [S, halfDim, 2]
+            let axisSin = sinInterleaved.reshaped([seqLen, axisDim])  // [S, axisDim]
+
+            cosComponents.append(axisCos)
+            sinComponents.append(axisSin)
+        }
+
+        // Now concatenate all ALREADY INTERLEAVED axis embeddings
+        let cosEmb = concatenated(cosComponents, axis: -1)  // [S, totalDims]
+        let sinEmb = concatenated(sinComponents, axis: -1)  // [S, totalDims]
+
+        return (cos: cosEmb, sin: sinEmb)
+    }
+
+    /// Apply rotary embeddings to query and key tensors
+    /// - Parameters:
+    ///   - q: Query tensor [B, H, S, D]
+    ///   - k: Key tensor [B, H, S, D]
+    ///   - cos: Cosine embeddings [S, D]
+    ///   - sin: Sine embeddings [S, D]
+    /// - Returns: Tuple of rotated (query, key)
+    public func apply(
+        query q: MLXArray,
+        key k: MLXArray,
+        cos cosEmb: MLXArray,
+        sin sinEmb: MLXArray
+    ) -> (query: MLXArray, key: MLXArray) {
+        let rotatedQ = applyRotary(q, cos: cosEmb, sin: sinEmb)
+        let rotatedK = applyRotary(k, cos: cosEmb, sin: sinEmb)
+        return (query: rotatedQ, key: rotatedK)
+    }
+
+    /// Apply rotary embedding to a single tensor
+    /// Uses the same approach as diffusers: treat consecutive pairs as (real, imag)
+    private func applyRotary(_ x: MLXArray, cos cosEmb: MLXArray, sin sinEmb: MLXArray) -> MLXArray {
+        // x shape: [B, H, S, D]
+        // cos/sin shape: [S, D]
+
+        // Try fused Metal kernel first (requires float32 cos/sin)
+        let cosFloat32 = cosEmb.dtype == .float32 ? cosEmb : cosEmb.asType(.float32)
+        let sinFloat32 = sinEmb.dtype == .float32 ? sinEmb : sinEmb.asType(.float32)
+
+        if let fused = Flux2FusedKernels.applyRotaryEmb(x, cos: cosFloat32, sin: sinFloat32) {
+            return fused
+        }
+
+        // Fall back to reference implementation
+        return applyRotaryReference(x, cos: cosEmb, sin: sinEmb)
+    }
+
+    /// Reference implementation of rotary embeddings (fallback when kernel unavailable)
+    private func applyRotaryReference(_ x: MLXArray, cos cosEmb: MLXArray, sin sinEmb: MLXArray) -> MLXArray {
+        // x shape: [B, H, S, D]
+        // cos/sin shape: [S, D]
+
+        // Reshape for broadcasting
+        let cosExpanded = cosEmb.expandedDimensions(axes: [0, 1])  // [1, 1, S, D]
+        let sinExpanded = sinEmb.expandedDimensions(axes: [0, 1])  // [1, 1, S, D]
+
+        // Diffusers approach: reshape x to [B, H, S, D//2, 2] and unbind to get x_real, x_imag
+        // x_real = x[..., 0::2], x_imag = x[..., 1::2]  (every other element)
+        // x_rotated = stack([-x_imag, x_real], dim=-1).flatten(3)
+
+        let batchSize = x.shape[0]
+        let numHeads = x.shape[1]
+        let seqLen = x.shape[2]
+        let dim = x.shape[3]
+        let halfDim = dim / 2
+
+        // Reshape to [B, H, S, D/2, 2]
+        let xReshaped = x.reshaped([batchSize, numHeads, seqLen, halfDim, 2])
+
+        // Get real and imag parts (consecutive pairs)
+        let xReal = xReshaped[0..., 0..., 0..., 0..., 0]  // [B, H, S, D/2]
+        let xImag = xReshaped[0..., 0..., 0..., 0..., 1]  // [B, H, S, D/2]
+
+        // Create rotated: [-x_imag, x_real] interleaved
+        let negXImag = -xImag
+        // Stack and flatten: [[-x_imag[0], x_real[0]], [-x_imag[1], x_real[1]], ...]
+        let xRotatedStacked = stacked([negXImag, xReal], axis: -1)  // [B, H, S, D/2, 2]
+        let xRotated = xRotatedStacked.reshaped([batchSize, numHeads, seqLen, dim])  // [B, H, S, D]
+
+        // Apply rotation: x * cos + x_rotated * sin
+        return x * cosExpanded + xRotated * sinExpanded
+    }
+}
+
+/// Generate position IDs for image latents using GPU-native operations
+/// - Parameters:
+///   - height: Image height in latent space
+///   - width: Image width in latent space
+/// - Returns: Position IDs [H*W, 4]
+public func generateImagePositionIDs(height: Int, width: Int) -> MLXArray {
+    // GPU-native implementation using meshgrid-like operations
+    // Create H and W coordinate grids
+    let hIndices = MLXArray.arange(height, dtype: .int32)  // [H]
+    let wIndices = MLXArray.arange(width, dtype: .int32)   // [W]
+
+    // Create meshgrid: expand H to [H, W] and W to [H, W] using broadcastTo
+    let hExpanded = hIndices.expandedDimensions(axis: 1)  // [H, 1]
+    let wExpanded = wIndices.expandedDimensions(axis: 0)  // [1, W]
+    let hGrid = MLX.broadcast(hExpanded, to: [height, width])  // [H, W]
+    let wGrid = MLX.broadcast(wExpanded, to: [height, width])  // [H, W]
+
+    // Flatten to [H*W]
+    let hFlat = hGrid.reshaped([height * width])
+    let wFlat = wGrid.reshaped([height * width])
+
+    // Create T=0 and L=0 columns
+    let zeros = MLXArray.zeros([height * width], dtype: .int32)
+
+    // Stack to create [H*W, 4] with format [T, H, W, L]
+    return MLX.stacked([zeros, hFlat, wFlat, zeros], axis: 1)
+}
+
+/// Generate position IDs for text sequence using GPU-native operations
+/// - Parameter length: Text sequence length
+/// - Returns: Position IDs [length, 4]
+public func generateTextPositionIDs(length: Int) -> MLXArray {
+    // GPU-native implementation
+    // Create L column as [0, 1, 2, ..., length-1]
+    let lIndices = MLXArray.arange(length, dtype: .int32)  // [length]
+
+    // Create T=0, H=0, W=0 columns
+    let zeros = MLXArray.zeros([length], dtype: .int32)
+
+    // Stack to create [length, 4] with format [T, H, W, L]
+    return MLX.stacked([zeros, zeros, zeros, lIndices], axis: 1)
+}
+
+/// Combine text and image position IDs for joint attention
+/// - Parameters:
+///   - textLength: Text sequence length
+///   - height: Image height in latent space
+///   - width: Image width in latent space
+/// - Returns: Combined position IDs [textLength + H*W, 4]
+public func generateCombinedPositionIDs(
+    textLength: Int,
+    height: Int,
+    width: Int
+) -> MLXArray {
+    let textIds = generateTextPositionIDs(length: textLength)
+    let imageIds = generateImagePositionIDs(height: height, width: width)
+    return concatenated([textIds, imageIds], axis: 0)
+}