From 6dd84dd01fa7bd7b944e0ae39e51b16d2256c761 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Fran=C3=A7ois=20Reboud?=
 <jean-francois.reboud@owkin.com>
Date: Fri, 28 Jun 2024 11:19:59 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(layer=5Fseq):=20QueryCausalSeq?=
 =?UTF-8?q?=20(#125)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CHANGELOG.md                                  |   1 +
 Sources/GrAIdient/LayerSeq/QuerySeq.swift     | 746 ++++++++++++++++++
 .../Metal/Kernel/LayerSeqFloat.metal          | 124 ---
 .../GrAIdient/Metal/Kernel/LayerSeqHalf.metal | 124 ---
 Sources/GrAIdient/Metal/Kernel/NLPFloat.metal | 497 ++++++++++++
 Sources/GrAIdient/Metal/Kernel/NLPHalf.metal  | 497 ++++++++++++
 Sources/GrAIdient/Metal/MetalConfig.swift     |  24 +-
 Sources/GrAIdient/Utils/Serialization.swift   |   1 +
 .../GrAIExamples/Base/python_lib/nlp/model.py |   6 +-
 Tests/GrAIExamples/NLPExample.swift           |  57 +-
 .../Base/InputSeq/EmbeddingSeqMSE1DCase.swift |   4 +-
 Tests/GrAITests/NLPTests.swift                | 358 +++++++++
 12 files changed, 2173 insertions(+), 266 deletions(-)
 create mode 100644 Sources/GrAIdient/Metal/Kernel/NLPFloat.metal
 create mode 100644 Sources/GrAIdient/Metal/Kernel/NLPHalf.metal

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f501fe0..84566f60 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+✨ **layer_seq:** QueryCausalSeq ([125](https://github.com/owkin/GrAIdient/pull/125))\
 ✨ **layer_seq:** RoPESeq ([124](https://github.com/owkin/GrAIdient/pull/124))\
 ✨ **layer_seq:** RMSNormSeq ([123](https://github.com/owkin/GrAIdient/pull/123))\
 ✨ **layer_seq:** EmbeddingSeq ([122](https://github.com/owkin/GrAIdient/pull/122))\
diff --git a/Sources/GrAIdient/LayerSeq/QuerySeq.swift b/Sources/GrAIdient/LayerSeq/QuerySeq.swift
index 180403cb..012fae53 100644
--- a/Sources/GrAIdient/LayerSeq/QuerySeq.swift
+++ b/Sources/GrAIdient/LayerSeq/QuerySeq.swift
@@ -996,3 +996,749 @@ public class QuerySelfSeq: LayerSeq
         }
     }
 }
+
+///
+/// Layer with a sequential shape neural structure.
+///
+/// This layer computes the causal attention scores between a query layer and a key layer.
+///
+public class QueryCausalSeq: LayerMergeSeq
+{
+    /// Number of heads (groups) of neurons for query.
+    let _nbHeadsQuery: Int
+    /// Number of heads (groups) of neurons for key.
+    let _nbHeadsKey: Int
+    
+    private enum Keys: String, CodingKey
+    {
+        case nbHeadsQuery
+        case nbHeadsKey
+    }
+    
+    ///
+    /// Create a layer with a sequential shape neural structure.
+    ///
+    /// - Parameters:
+    ///     - query: Previous layer containing the query to look for.
+    ///     - key: Previous layer containing the keys of reference.
+    ///     - nbHeadsQuery: Number of heads (groups) of neurons for query.
+    ///     - nbHeadsKey: Number of heads (groups) of neurons for key.
+    ///     - params: Contextual parameters linking to the model.
+    ///
+    public init(query: LayerSeq, key: LayerSeq, 
+                nbHeadsQuery: Int, nbHeadsKey: Int,
+                params: GrAI.Model.Params) throws
+    {
+        if query.nbNeurons % nbHeadsQuery != 0
+        {
+            throw LayerError.Init(message:
+                "`nbNeurons` (\(query.nbNeurons)) " +
+                "should be a multiple of `nbHeadsQuery` (\(nbHeadsQuery))."
+            )
+        }
+        if key.nbNeurons % nbHeadsKey != 0
+        {
+            throw LayerError.Init(message:
+                "`nbNeurons` (\(key.nbNeurons)) " +
+                "should be a multiple of `nbHeadsKey` (\(nbHeadsKey))."
+            )
+        }
+        if nbHeadsQuery % nbHeadsKey != 0
+        {
+            throw LayerError.Init(message:
+                "`nbHeadsQuery` should be a multiple of `nbHeadsKey`"
+            )
+        }
+        if query.nbNeurons / nbHeadsQuery != key.nbNeurons / nbHeadsKey
+        {
+            throw LayerError.Init(message:
+                "`query` and `key` should should have same hidden dimension."
+            )
+        }
+        if query.sequence != key.sequence
+        {
+            throw LayerError.Init(message: "Layer structure error.")
+        }
+        
+        _nbHeadsQuery = nbHeadsQuery
+        _nbHeadsKey = nbHeadsKey
+        
+        super.init(layersPrev: [query, key],
+                   sequence: query.sequence,
+                   nbNeurons: query.sequence * nbHeadsQuery,
+                   params: params)
+    }
+    
+    ///
+    /// Decode from the disk.
+    ///
+    /// Throw an error if reading from the decoder fails, or
+    /// if the data read is corrupted or otherwise invalid.
+    ///
+    /// - Parameter decoder: The decoder to read data from.
+    ///
+    public required init(from decoder: Decoder) throws
+    {
+        let values = try decoder.container(keyedBy: Keys.self)
+        _nbHeadsQuery = try values.decode(Int.self, forKey: Keys.nbHeadsQuery)
+        _nbHeadsKey = try values.decode(Int.self, forKey: Keys.nbHeadsKey)
+        try super.init(from: decoder)
+    }
+    
+    ///
+    /// Encode to the disk.
+    ///
+    /// If the value fails to encode anything, `encoder` will encode an empty
+    /// keyed container in its place.
+    ///
+    /// Throw an error if any values are invalid for the given
+    /// encoder's format.
+    ///
+    /// - Parameter encoder: The encoder to write data to.
+    ///
+    public override func encode(to encoder: Encoder) throws
+    {
+        var container = encoder.container(keyedBy: Keys.self)
+        try container.encode(_nbHeadsQuery, forKey: Keys.nbHeadsQuery)
+        try container.encode(_nbHeadsKey, forKey: Keys.nbHeadsKey)
+        try super.encode(to: encoder)
+    }
+    
+    ///
+    /// Create a layer with same values as this.
+    ///
+    /// - Parameters:
+    ///     - mapping: Dictionary allowing to find the layer associated to some id.
+    ///     This dictionary is particularly useful when the different layers cannot access
+    ///     their `layerPrev`.
+    ///     - inPlace: Whether hard resources should be copied as is.
+    ///
+    /// - Returns: A new layer. When `inPlace` is false, `initKernel` is
+    /// necessary in order to recreate hard resources.
+    ///
+    public override func copy(
+        mapping: Dictionary<Int, Layer>,
+        inPlace: Bool) -> Layer
+    {
+        let context = ModelContext(name: "", curID: 0)
+        let params = GrAI.Model.Params(context: context)
+        params.context.curID = id
+        
+        var layersPrev = [LayerSeq]()
+        for idPrev in _idsPrev
+        {
+            layersPrev.append(mapping[idPrev] as! LayerSeq)
+        }
+        
+        let layer = try! QueryCausalSeq(
+            query: layersPrev[0], key: layersPrev[1],
+            nbHeadsQuery: _nbHeadsQuery,
+            nbHeadsKey: _nbHeadsKey,
+            params: params
+        )
+        return layer
+    }
+    
+    ///
+    /// Initialize state resources in the CPU execution context.
+    ///
+    /// We initialize the neurons' state (forward and backward).
+    ///
+    public override func checkStateCPU(batchSize: Int) throws
+    {
+        if neurons == nil
+        {
+            try super.checkStateCPU(batchSize: batchSize)
+            _encodeCausalityCPU()
+        }
+        else
+        {
+            try super.checkStateCPU(batchSize: batchSize)
+        }
+    }
+    
+    /// Update causality scores in the CPU execution context.
+    private func _encodeCausalityCPU()
+    {
+        for elem in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in seqQ+1..<sequence
+        {
+            neurons.get(seqQ, seqK + headQuery * sequence)!.v[elem].out = -1e9
+        }}}}
+    }
+    
+    ///
+    /// Initialize state resources in the GPU execution context.
+    ///
+    /// We initialize the neurons' forward state.
+    ///
+    public override func checkStateForwardGPU(batchSize: Int) throws
+    {
+        if outs == nil
+        {
+            try super.checkStateForwardGPU(batchSize: batchSize)
+            _encodeCausalityGPU()
+        }
+        else
+        {
+            try super.checkStateForwardGPU(batchSize: batchSize)
+        }
+    }
+    
+    /// Update causality scores in the GPU execution context.
+    private func _encodeCausalityGPU()
+    {
+        let pNbHeadsQuery: [UInt32] = [UInt32(_nbHeadsQuery)]
+        let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
+        let pNbBatch: [UInt32] = [UInt32(batchSize)]
+        let pSequence: [UInt32] = [UInt32(sequence)]
+        
+        let command = MetalKernel.get.createCommand(
+            "encodeCausality", deviceID: deviceID
+        )
+        command.setBytes(pNbHeadsQuery, atIndex: 0)
+        command.setBytes(pNbNeurons, atIndex: 1)
+        command.setBytes(pNbBatch, atIndex: 2)
+        command.setBytes(pSequence, atIndex: 3)
+        command.setBuffer(outs.metal, atIndex: 4)
+        
+        command.dispatchThreads(
+            width: nbNeurons,
+            height: batchSize * sequence
+        )
+        command.enqueue()
+    }
+    
+    ///
+    /// Apply the forward pass of the Gradient Checking in CPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGCCPU() throws
+    {
+        try checkStateCPU(batchSize: batchSize)
+        
+        let (nbSameElems, layersIndex, nbElems) = getMergedGraph()
+        
+        var nbGC = nbSameElems
+        for nbElemsTmp in nbElems
+        {
+            nbGC += nbElemsTmp
+        }
+        
+        for seqQ in 0..<sequence {
+        for seqK in 0..<nbNeurons
+        {
+            neurons.get(seqQ, seqK)!.initGC(batchSize: batchSize, nbGC: nbGC)
+        }}
+        
+        let query = (_layersPrev[0] as! LayerSeq).neurons!
+        let key = (_layersPrev[1] as! LayerSeq).neurons!
+        let size = (_layersPrev[0] as! LayerSeq).nbNeurons / _nbHeadsQuery
+        
+        for batch in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in 0..<sequence {
+        for elem in 0..<nbSameElems
+        {
+            if seqK <= seqQ
+            {
+                let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                    headQuery : headQuery / _nbHeadsKey
+                var sum = 0.0
+                
+                for j in 0..<size
+                {
+                    let depthPrevKey = j + headKey * size
+                    let depthPrevQuery = j + headQuery * size
+                    
+                    let queryTmp = query.get(
+                        seqQ, depthPrevQuery
+                    )!.gc[batch][elem].out
+                    let keyTmp = key.get(
+                        seqK, depthPrevKey
+                    )!.gc[batch][elem].out
+                    
+                    sum += queryTmp * keyTmp
+                }
+                
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][elem].out = sum / sqrt(Double(size))
+            }
+            else
+            {
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][elem].out = -1e9
+            }
+        }}}}}
+        
+        for batch in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in 0..<sequence {
+        var offset = nbSameElems
+        var nbLastElems = [Int](repeating: nbSameElems,
+                                count: _layersPrev.count)
+        for (index, nbElemsTmp) in zip(layersIndex, nbElems) {
+        for elem in 0..<nbElemsTmp
+        {
+            if seqK <= seqQ
+            {
+                let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                    headQuery : headQuery / _nbHeadsKey
+                var sum = 0.0
+                
+                for j in 0..<size
+                {
+                    let depthPrevKey = j + headKey * size
+                    let depthPrevQuery = j + headQuery * size
+                    
+                    let queryTmp: Double
+                    let keyTmp: Double
+                    
+                    if index == 0
+                    {
+                        queryTmp = query.get(
+                            seqQ, depthPrevQuery
+                        )!.gc[batch][nbLastElems[index]+elem].out
+                        keyTmp = key.get(seqK, depthPrevKey)!.v[batch].out
+                    }
+                    else
+                    {
+                        queryTmp = query.get(seqQ, depthPrevQuery)!.v[batch].out
+                        keyTmp = key.get(
+                            seqK, depthPrevKey
+                        )!.gc[batch][nbLastElems[index]+elem].out
+                    }
+                    
+                    sum += queryTmp * keyTmp
+                }
+                
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][offset+elem].out = sum / sqrt(Double(size))
+            }
+            else
+            {
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][offset+elem].out = -1e9
+            }
+        }
+        
+        offset += nbElemsTmp
+        nbLastElems[index] += nbElemsTmp
+        }}}}}
+    }
+    
+    ///
+    /// Apply the forward pass of the Gradient Checking in GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGCGPU() throws
+    {
+        try checkStateCPU(batchSize: batchSize)
+        
+        let (nbSameElems, layersIndex, nbElems) = getMergedGraph()
+        
+        var nbGC = nbSameElems
+        for nbElemsTmp in nbElems
+        {
+            nbGC += nbElemsTmp
+        }
+        
+        for seqQ in 0..<sequence {
+        for seqK in 0..<nbNeurons
+        {
+            neurons.get(seqQ, seqK)!.initGC(batchSize: batchSize, nbGC: nbGC)
+        }}
+        
+        let query = (_layersPrev[0] as! LayerSeq).neurons!
+        let key = (_layersPrev[1] as! LayerSeq).neurons!
+        let nbNeuronsPrevQuery = (_layersPrev[0] as! LayerSeq).nbNeurons
+        let nbNeuronsPrevKey = (_layersPrev[1] as! LayerSeq).nbNeurons
+        let size = (_layersPrev[0] as! LayerSeq).nbNeurons / _nbHeadsQuery
+        
+        for batch in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in 0..<sequence {
+        for elem in 0..<nbSameElems
+        {
+            if seqK <= seqQ
+            {
+                let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                    headQuery : headQuery / _nbHeadsKey
+                var sum = 0.0
+                
+                for j in 0..<size
+                {
+                    let depthPrevKey = j + headKey * size
+                    let depthPrevQuery = j + headQuery * size
+                    
+                    let queryTmp = query.get(
+                        seqQ, depthPrevQuery
+                    )!.gc[batch][elem].out
+                    let keyTmp = key.get(
+                        seqK, depthPrevKey
+                    )!.gc[batch][elem].out
+                    
+                    sum += queryTmp * keyTmp
+                }
+                
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][elem].out = sum / sqrt(Double(size))
+            }
+            else
+            {
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][elem].out = -1e9
+            }
+        }}}}}
+        
+        let queryBuffer = (_layersPrev[0] as! LayerSeq).outs.download()
+        let keyBuffer = (_layersPrev[1] as! LayerSeq).outs.download()
+        
+        for batch in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in 0..<sequence {
+        var offset = nbSameElems
+        var nbLastElems = [Int](repeating: nbSameElems,
+                                count: _layersPrev.count)
+        for (index, nbElemsTmp) in zip(layersIndex, nbElems) {
+        for elem in 0..<nbElemsTmp
+        {
+            if seqK <= seqQ
+            {
+                let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                    headQuery : headQuery / _nbHeadsKey
+                var sum = 0.0
+                
+                for j in 0..<size
+                {
+                    let depthPrevKey = j + headKey * size
+                    let depthPrevQuery = j + headQuery * size
+                    
+                    let queryTmp: Double
+                    let keyTmp: Double
+                    
+                    if index == 0
+                    {
+                        queryTmp = query.get(
+                            seqQ, depthPrevQuery
+                        )!.gc[batch][nbLastElems[index]+elem].out
+                        
+                        let offsetTmp = depthPrevKey + 
+                            nbNeuronsPrevKey * seqK +
+                            sequence * nbNeuronsPrevKey * batch
+                        
+                        keyTmp = Double(keyBuffer[offsetTmp])
+                    }
+                    else
+                    {
+                        let offsetTmp = depthPrevQuery + 
+                            nbNeuronsPrevQuery * seqQ +
+                            sequence * nbNeuronsPrevQuery * batch
+                        
+                        queryTmp = Double(queryBuffer[offsetTmp])
+                        
+                        keyTmp = key.get(
+                            seqK, depthPrevKey
+                        )!.gc[batch][nbLastElems[index]+elem].out
+                    }
+                    
+                    sum += queryTmp * keyTmp
+                }
+                
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][offset+elem].out = sum / sqrt(Double(size))
+            }
+            else
+            {
+                neurons.get(
+                    seqQ, seqK + headQuery * sequence
+                )!.gc[batch][offset+elem].out = -1e9
+            }
+        }
+        
+        offset += nbElemsTmp
+        nbLastElems[index] += nbElemsTmp
+        }}}}}
+    }
+    
+    ///
+    /// Apply the forward pass in the CPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardCPU() throws
+    {
+        try checkStateCPU(batchSize: batchSize)
+        
+        let query = (_layersPrev[0] as! LayerSeq).neurons!
+        let key = (_layersPrev[1] as! LayerSeq).neurons!
+        let size = (_layersPrev[0] as! LayerSeq).nbNeurons / _nbHeadsQuery
+        
+        for elem in 0..<batchSize {
+        for headQuery in 0..<_nbHeadsQuery {
+        for seqQ in 0..<sequence {
+        for seqK in 0...seqQ
+        {
+            let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                headQuery : headQuery / _nbHeadsKey
+            var sum = 0.0
+            
+            for j in 0..<size
+            {
+                let depthPrevKey = j + headKey * size
+                let depthPrevQuery = j + headQuery * size
+                
+                let queryTmp = query.get(seqQ, depthPrevQuery)!.v[elem].out
+                let keyTmp = key.get(seqK, depthPrevKey)!.v[elem].out
+                
+                sum += queryTmp * keyTmp
+            }
+            
+            neurons.get(seqQ, seqK + headQuery * sequence)!.v[elem].out =
+                sum / sqrt(Double(size))
+        }}}}
+    }
+    
+    ///
+    /// Apply the forward pass in the GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGPU() throws
+    {
+        try checkStateForwardGPU(batchSize: batchSize)
+        
+        let query = _layersPrev[0] as! LayerSeq
+        let key = _layersPrev[1] as! LayerSeq
+        let nbNeuronsPrevQuery = query.nbNeurons
+        let nbNeuronsPrevKey = key.nbNeurons
+        
+        let pNbHeadsQuery: [UInt32] = [UInt32(_nbHeadsQuery)]
+        let pNbHeadsKey: [UInt32] = [UInt32(_nbHeadsKey)]
+        let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
+        let pNbNeuronsPrevQuery: [UInt32] = [UInt32(nbNeuronsPrevQuery)]
+        let pNbNeuronsPrevKey: [UInt32] = [UInt32(nbNeuronsPrevKey)]
+        let pNbBatch: [UInt32] = [UInt32(batchSize)]
+        let pSequence: [UInt32] = [UInt32(sequence)]
+        
+        let kernel = (nbNeuronsPrevQuery / _nbHeadsQuery) % 4 == 0 ?
+            "queryCausalSeq4Forward" : "queryCausalSeqForward"
+        let command = MetalKernel.get.createCommand(
+            kernel, deviceID: deviceID
+        )
+        command.setBuffer(query.outs.metal, atIndex: 0)
+        command.setBuffer(key.outs.metal, atIndex: 1)
+        command.setBytes(pNbHeadsQuery, atIndex: 2)
+        command.setBytes(pNbHeadsKey, atIndex: 3)
+        command.setBytes(pNbNeurons, atIndex: 4)
+        command.setBytes(pNbNeuronsPrevQuery, atIndex: 5)
+        command.setBytes(pNbNeuronsPrevKey, atIndex: 6)
+        command.setBytes(pNbBatch, atIndex: 7)
+        command.setBytes(pSequence, atIndex: 8)
+        command.setBuffer(outs.metal, atIndex: 9)
+        
+        command.dispatchThreads(
+            width: nbNeurons,
+            height: batchSize * sequence
+        )
+        command.enqueue()
+    }
+    
+    /// Apply the backward pass in the CPU execution context.
+    public override func backwardCPU()
+    {
+        if !mustComputeBackward
+        {
+            return
+        }
+        
+        let query = (_layersPrev[0] as! LayerSeq).neurons!
+        let key = (_layersPrev[1] as! LayerSeq).neurons!
+        let size = (_layersPrev[0] as! LayerSeq).nbNeurons / _nbHeadsQuery
+        
+        if _layersPrev[0].computeDelta
+        {
+            for elem in 0..<batchSize {
+            for headQuery in 0..<_nbHeadsQuery {
+            let headKey = _nbHeadsQuery == _nbHeadsKey ?
+                headQuery : headQuery / _nbHeadsKey
+            for seqQ in 0..<sequence {
+            for j in 0..<size
+            {
+                let depthPrevKey = j + headKey * size
+                let depthPrevQuery = j + headQuery * size
+                
+                var sum = 0.0
+                for seqK in 0...seqQ
+                {
+                    let deltaCur = neurons
+                        .get(seqQ, seqK + headQuery * sequence)!.v[elem].delta
+                    let keyTmp = key.get(seqK, depthPrevKey)!.v[elem].out
+                    
+                    sum += deltaCur * keyTmp
+                }
+                
+                if _layersPrev[0].dirty
+                {
+                    query.get(seqQ, depthPrevQuery)!.v[elem].delta =
+                        sum / sqrt(Double(size))
+                }
+                else
+                {
+                    query.get(seqQ, depthPrevQuery)!.v[elem].delta +=
+                        sum / sqrt(Double(size))
+                }
+            }}}}
+        }
+        if _layersPrev[1].computeDelta
+        {
+            let nbBlocksHead = _nbHeadsQuery == _nbHeadsKey ?
+                1 : _nbHeadsQuery / _nbHeadsKey
+            
+            for elem in 0..<batchSize {
+            for headKey in 0..<_nbHeadsKey {
+            for seqK in 0..<sequence {
+            for j in 0..<size
+            {
+                let depthPrevKey = j + headKey * size
+                
+                var sum = 0.0
+                for blockHead in 0..<nbBlocksHead
+                {
+                    let headQuery = blockHead + nbBlocksHead * headKey
+                    let depthPrevQuery = j + headQuery * size
+                    
+                    for seqQ in seqK..<sequence
+                    {
+                        let deltaCur = neurons.get(
+                            seqQ, seqK + headQuery * sequence
+                        )!.v[elem].delta
+                        let queryTmp = query.get(
+                            seqQ, depthPrevQuery
+                        )!.v[elem].out
+                        
+                        sum += deltaCur * queryTmp
+                    }
+                }
+                
+                if _layersPrev[1].dirty
+                {
+                    key.get(seqK, depthPrevKey)!.v[elem].delta =
+                        sum / sqrt(Double(size))
+                }
+                else
+                {
+                    key.get(seqK, depthPrevKey)!.v[elem].delta +=
+                        sum / sqrt(Double(size))
+                }
+            }}}}
+        }
+        propagateDirty()
+    }
+    
+    ///
+    /// Apply the backward pass in the GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func backwardGPU() throws
+    {
+        if !mustComputeBackward
+        {
+            return
+        }
+        
+        let query = _layersPrev[0] as! LayerSeq
+        let key = _layersPrev[1] as! LayerSeq
+        let nbNeuronsPrevQuery = query.nbNeurons
+        let nbNeuronsPrevKey = key.nbNeurons
+        
+        let pNbHeadsQuery: [UInt32] = [UInt32(_nbHeadsQuery)]
+        let pNbHeadsKey: [UInt32] = [UInt32(_nbHeadsKey)]
+        let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
+        let pNbNeuronsPrevQuery: [UInt32] = [UInt32(nbNeuronsPrevQuery)]
+        let pNbNeuronsPrevKey: [UInt32] = [UInt32(nbNeuronsPrevKey)]
+        let pNbBatch: [UInt32] = [UInt32(batchSize)]
+        let pSequence: [UInt32] = [UInt32(sequence)]
+        
+        let metalKernel = MetalKernel.get
+        var command: MetalCommand
+        
+        if query.computeDelta
+        {
+            try query.checkStateBackwardGPU(batchSize: batchSize)
+            
+            let pDirty: [UInt32] = query.dirty ? [1] : [0]
+            
+            let kernel = (nbNeuronsPrevQuery / _nbHeadsQuery) % 4 == 0 ?
+                "queryCausalQuerySeq4Backward" : "queryCausalQuerySeqBackward"
+            let coeff = (nbNeuronsPrevQuery / _nbHeadsQuery) % 4 == 0 ? 4 : 1
+            command = metalKernel.createCommand(
+                kernel, deviceID: deviceID
+            )
+            command.setBuffer(delta.metal, atIndex: 0)
+            command.setBuffer(key.outs.metal, atIndex: 1)
+            command.setBytes(pNbHeadsQuery, atIndex: 2)
+            command.setBytes(pNbHeadsKey, atIndex: 3)
+            command.setBytes(pNbNeurons, atIndex: 4)
+            command.setBytes(pNbNeuronsPrevQuery, atIndex: 5)
+            command.setBytes(pNbNeuronsPrevKey, atIndex: 6)
+            command.setBytes(pNbBatch, atIndex: 7)
+            command.setBytes(pSequence, atIndex: 8)
+            command.setBytes(pDirty, atIndex: 9)
+            command.setBuffer(query.delta.metal, atIndex: 10)
+            
+            command.dispatchThreads(
+                width: nbNeuronsPrevQuery / coeff,
+                height: batchSize * sequence
+            )
+            command.enqueue()
+        }
+        if key.computeDelta
+        {
+            try key.checkStateBackwardGPU(batchSize: batchSize)
+            
+            let pDirty: [UInt32] = key.dirty ? [1] : [0]
+            
+            let kernel = (nbNeuronsPrevKey / _nbHeadsKey) % 4 == 0 ?
+                "queryCausalKeySeq4Backward" : "queryCausalKeySeqBackward"
+            let coeff = (nbNeuronsPrevKey / _nbHeadsKey) % 4 == 0 ? 4 : 1
+            command = metalKernel.createCommand(
+                kernel, deviceID: deviceID
+            )
+            command.setBuffer(delta.metal, atIndex: 0)
+            command.setBuffer(query.outs.metal, atIndex: 1)
+            command.setBytes(pNbHeadsQuery, atIndex: 2)
+            command.setBytes(pNbHeadsKey, atIndex: 3)
+            command.setBytes(pNbNeurons, atIndex: 4)
+            command.setBytes(pNbNeuronsPrevQuery, atIndex: 5)
+            command.setBytes(pNbNeuronsPrevKey, atIndex: 6)
+            command.setBytes(pNbBatch, atIndex: 7)
+            command.setBytes(pSequence, atIndex: 8)
+            command.setBytes(pDirty, atIndex: 9)
+            command.setBuffer(key.delta.metal, atIndex: 10)
+            
+            command.dispatchThreads(
+                width: nbNeuronsPrevKey / coeff,
+                height: batchSize * sequence
+            )
+            command.enqueue()
+        }
+        propagateDirty()
+    }
+}
diff --git a/Sources/GrAIdient/Metal/Kernel/LayerSeqFloat.metal b/Sources/GrAIdient/Metal/Kernel/LayerSeqFloat.metal
index f8694e96..787c985f 100644
--- a/Sources/GrAIdient/Metal/Kernel/LayerSeqFloat.metal
+++ b/Sources/GrAIdient/Metal/Kernel/LayerSeqFloat.metal
@@ -2743,127 +2743,3 @@ kernel void layerCAMSeqForwardFloat(
     uint offset = seq + sequence * elem;
     outs[offset] = sum;
 }
-
-kernel void createRoPESeqMatrixFloat(
-    constant int * seqPositions,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & sequence,
-    device float * rotationMatrix,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint block = id[0];
-    uint seq = id[1];
-    
-    if (block >= nbBlocks || seq >= sequence)
-    {
-        return ;
-    }
-    
-    float position = (float)seqPositions[seq];
-    float theta = pow(
-        10000.0,
-        -2.0 * (float)block / (float)size
-    );
-    float mTheta = position * theta;
-    float cosVal = cos(mTheta);
-    float sinVal = sin(mTheta);
-    
-    uint offset = 2 * block + seq * size;
-    rotationMatrix[offset] = cosVal;
-    rotationMatrix[1 + offset] = sinVal;
-}
-
-kernel void RoPESeqForwardFloat(
-    const device float * outsPrev,
-    const device float * rotationMatrix,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & nbBatch,
-    constant uint & sequence,
-    device float * outs,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint head = id[0] / nbBlocks;
-    uint block = id[0] % nbBlocks;
-    uint elem = id[1] / sequence;
-    uint seq = id[1] % sequence;
-    
-    if (head >= nbHeads || block >= nbBlocks ||
-        elem >= nbBatch || seq >= sequence)
-    {
-        return ;
-    }
-    
-    uint offset1 = 2 * block + seq * size;
-    uint offset2 = 2 * block + head * size +
-        nbNeurons * seq + sequence * nbNeurons * elem;
-    
-    float cosVal = rotationMatrix[offset1];
-    float sinVal = rotationMatrix[1 + offset1];
-    
-    float in1 = outsPrev[offset2];
-    float in2 = outsPrev[1 + offset2];
-    
-    float out1 = in1 * cosVal - in2 * sinVal;
-    float out2 = in1 * sinVal + in2 * cosVal;
-    
-    outs[offset2] = out1;
-    outs[1 + offset2] = out2;
-}
-
-kernel void RoPESeqSeqBackwardFloat(
-    const device float * delta,
-    const device float * rotationMatrix,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & nbBatch,
-    constant uint & sequence,
-    constant uint & dirty,
-    device float * deltaPrev,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint head = id[0] / nbBlocks;
-    uint block = id[0] % nbBlocks;
-    uint elem = id[1] / sequence;
-    uint seq = id[1] % sequence;
-    
-    if (head >= nbHeads || block >= nbBlocks ||
-        elem >= nbBatch || seq >= sequence)
-    {
-        return ;
-    }
-    
-    uint offset1 = 2 * block + seq * size;
-    uint offset2 = 2 * block + head * size +
-        nbNeurons * seq + sequence * nbNeurons * elem;
-    
-    float cosVal = rotationMatrix[offset1];
-    float sinVal = rotationMatrix[1 + offset1];
-    
-    float out1 = delta[offset2];
-    float out2 = delta[1 + offset2];
-    
-    float in1 = out1 * cosVal + out2 * sinVal;
-    float in2 = -out1 * sinVal + out2 * cosVal;
-    
-    if (dirty)
-    {
-        deltaPrev[offset2] = in1;
-        deltaPrev[1 + offset2] = in2;
-    }
-    else
-    {
-        deltaPrev[offset2] += in1;
-        deltaPrev[1 + offset2] += in2;
-    }
-}
diff --git a/Sources/GrAIdient/Metal/Kernel/LayerSeqHalf.metal b/Sources/GrAIdient/Metal/Kernel/LayerSeqHalf.metal
index 80f86c7d..21a2a7be 100644
--- a/Sources/GrAIdient/Metal/Kernel/LayerSeqHalf.metal
+++ b/Sources/GrAIdient/Metal/Kernel/LayerSeqHalf.metal
@@ -2743,127 +2743,3 @@ kernel void layerCAMSeqForwardHalf(
     uint offset = seq + sequence * elem;
     outs[offset] = sum;
 }
-
-kernel void createRoPESeqMatrixHalf(
-    constant int * seqPositions,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & sequence,
-    device half * rotationMatrix,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint block = id[0];
-    uint seq = id[1];
-    
-    if (block >= nbBlocks || seq >= sequence)
-    {
-        return ;
-    }
-    
-    float position = (float)seqPositions[seq];
-    float theta = pow(
-        10000.0,
-        -2.0 * (float)block / (float)size
-    );
-    float mTheta = position * theta;
-    float cosVal = cos(mTheta);
-    float sinVal = sin(mTheta);
-    
-    uint offset = 2 * block + seq * size;
-    rotationMatrix[offset] = cosVal;
-    rotationMatrix[1 + offset] = sinVal;
-}
-
-kernel void RoPESeqForwardHalf(
-    const device half * outsPrev,
-    const device half * rotationMatrix,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & nbBatch,
-    constant uint & sequence,
-    device half * outs,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint head = id[0] / nbBlocks;
-    uint block = id[0] % nbBlocks;
-    uint elem = id[1] / sequence;
-    uint seq = id[1] % sequence;
-    
-    if (head >= nbHeads || block >= nbBlocks ||
-        elem >= nbBatch || seq >= sequence)
-    {
-        return ;
-    }
-    
-    uint offset1 = 2 * block + seq * size;
-    uint offset2 = 2 * block + head * size +
-        nbNeurons * seq + sequence * nbNeurons * elem;
-    
-    half cosVal = rotationMatrix[offset1];
-    half sinVal = rotationMatrix[1 + offset1];
-    
-    half in1 = outsPrev[offset2];
-    half in2 = outsPrev[1 + offset2];
-    
-    half out1 = in1 * cosVal - in2 * sinVal;
-    half out2 = in1 * sinVal + in2 * cosVal;
-    
-    outs[offset2] = out1;
-    outs[1 + offset2] = out2;
-}
-
-kernel void RoPESeqSeqBackwardHalf(
-    const device half * delta,
-    const device half * rotationMatrix,
-    constant uint & nbHeads,
-    constant uint & nbNeurons,
-    constant uint & nbBatch,
-    constant uint & sequence,
-    constant uint & dirty,
-    device half * deltaPrev,
-    uint2 id [[ thread_position_in_grid ]])
-{
-    uint size = nbNeurons / nbHeads;
-    uint nbBlocks = size / 2;
-    
-    uint head = id[0] / nbBlocks;
-    uint block = id[0] % nbBlocks;
-    uint elem = id[1] / sequence;
-    uint seq = id[1] % sequence;
-    
-    if (head >= nbHeads || block >= nbBlocks ||
-        elem >= nbBatch || seq >= sequence)
-    {
-        return ;
-    }
-    
-    uint offset1 = 2 * block + seq * size;
-    uint offset2 = 2 * block + head * size +
-        nbNeurons * seq + sequence * nbNeurons * elem;
-    
-    half cosVal = rotationMatrix[offset1];
-    half sinVal = rotationMatrix[1 + offset1];
-    
-    half out1 = delta[offset2];
-    half out2 = delta[1 + offset2];
-    
-    half in1 = out1 * cosVal + out2 * sinVal;
-    half in2 = -out1 * sinVal + out2 * cosVal;
-    
-    if (dirty)
-    {
-        deltaPrev[offset2] = in1;
-        deltaPrev[1 + offset2] = in2;
-    }
-    else
-    {
-        deltaPrev[offset2] += in1;
-        deltaPrev[1 + offset2] += in2;
-    }
-}
diff --git a/Sources/GrAIdient/Metal/Kernel/NLPFloat.metal b/Sources/GrAIdient/Metal/Kernel/NLPFloat.metal
new file mode 100644
index 00000000..89ad05c7
--- /dev/null
+++ b/Sources/GrAIdient/Metal/Kernel/NLPFloat.metal
@@ -0,0 +1,497 @@
+//
+// NLPFloat.metal
+// GrAIdient
+//
+// Created by Jean-François Reboud on 25/06/2024.
+//
+
+#include <metal_stdlib>
+using namespace metal;
+
+kernel void createRoPESeqMatrixFloat(
+    constant int * seqPositions,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & sequence,
+    device float * rotationMatrix,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint block = id[0];
+    uint seq = id[1];
+    
+    if (block >= nbBlocks || seq >= sequence)
+    {
+        return ;
+    }
+    
+    float position = (float)seqPositions[seq];
+    float theta = pow(
+        10000.0,
+        -2.0 * (float)block / (float)size
+    );
+    float mTheta = position * theta;
+    float cosVal = cos(mTheta);
+    float sinVal = sin(mTheta);
+    
+    uint offset = 2 * block + seq * size;
+    rotationMatrix[offset] = cosVal;
+    rotationMatrix[1 + offset] = sinVal;
+}
+
+kernel void RoPESeqForwardFloat(
+    const device float * outsPrev,
+    const device float * rotationMatrix,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device float * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint head = id[0] / nbBlocks;
+    uint block = id[0] % nbBlocks;
+    uint elem = id[1] / sequence;
+    uint seq = id[1] % sequence;
+    
+    if (head >= nbHeads || block >= nbBlocks ||
+        elem >= nbBatch || seq >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset1 = 2 * block + seq * size;
+    uint offset2 = 2 * block + head * size +
+        nbNeurons * seq + sequence * nbNeurons * elem;
+    
+    float cosVal = rotationMatrix[offset1];
+    float sinVal = rotationMatrix[1 + offset1];
+    
+    float in1 = outsPrev[offset2];
+    float in2 = outsPrev[1 + offset2];
+    
+    float out1 = in1 * cosVal - in2 * sinVal;
+    float out2 = in1 * sinVal + in2 * cosVal;
+    
+    outs[offset2] = out1;
+    outs[1 + offset2] = out2;
+}
+
+kernel void RoPESeqSeqBackwardFloat(
+    const device float * delta,
+    const device float * rotationMatrix,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device float * deltaPrev,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint head = id[0] / nbBlocks;
+    uint block = id[0] % nbBlocks;
+    uint elem = id[1] / sequence;
+    uint seq = id[1] % sequence;
+    
+    if (head >= nbHeads || block >= nbBlocks ||
+        elem >= nbBatch || seq >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset1 = 2 * block + seq * size;
+    uint offset2 = 2 * block + head * size +
+        nbNeurons * seq + sequence * nbNeurons * elem;
+    
+    float cosVal = rotationMatrix[offset1];
+    float sinVal = rotationMatrix[1 + offset1];
+    
+    float out1 = delta[offset2];
+    float out2 = delta[1 + offset2];
+    
+    float in1 = out1 * cosVal + out2 * sinVal;
+    float in2 = -out1 * sinVal + out2 * cosVal;
+    
+    if (dirty)
+    {
+        deltaPrev[offset2] = in1;
+        deltaPrev[1 + offset2] = in2;
+    }
+    else
+    {
+        deltaPrev[offset2] += in1;
+        deltaPrev[1 + offset2] += in2;
+    }
+}
+
+kernel void encodeCausalityFloat(
+    constant uint & nbHeadsQuery,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device float * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery ||
+        seqK >= sequence || seqK <= seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = -1e9;
+}
+
+kernel void queryCausalSeqForwardFloat(
+    const device float * query,
+    const device float * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device float * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || seqK > seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    float tmp = 0.0;
+    
+    for (uint j=0; j<size; j++)
+    {
+        uint depthPrevKey = j + headKey * size;
+        uint depthPrevQuery = j + headQuery * size;
+        
+        uint offsetQuery = depthPrevQuery +
+            nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem;
+        uint offsetKey = depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+        
+        tmp += query[offsetQuery] * key[offsetKey];
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = tmp;
+}
+
+kernel void queryCausalSeq4ForwardFloat(
+    const device float4 * query,
+    const device float4 * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device float * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || seqK > seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    float4 tmp = 0.0;
+    
+    for (uint j=0; j<size/4; j++)
+    {
+        uint depthPrevKey = j * 4 + headKey * size;
+        uint depthPrevQuery = j * 4 + headQuery * size;
+        
+        uint offsetQuery = (depthPrevQuery +
+            nbNeuronsPrevQuery * seqQ +
+            sequence * nbNeuronsPrevQuery * elem) / 4;
+        uint offsetKey = (depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+        
+        tmp += query[offsetQuery] * key[offsetKey];
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = tmp[0] + tmp[1] + tmp[2] + tmp[3];
+}
+
+kernel void queryCausalQuerySeqBackwardFloat(
+    const device float * delta,
+    const device float * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device float * query,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / size;
+    uint j = id[0] % size;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || j >= size ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    uint depthPrevKey = j + headKey * size;
+    uint depthPrevQuery = j + headQuery * size;
+    
+    float tmp = 0.0;
+    for (uint seqK=0; seqK<=seqQ; seqK++)
+    {
+        uint offset = seqK + headQuery * sequence +
+            nbNeurons * seqQ + sequence * nbNeurons * elem;
+        uint offsetKey = depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+        
+        tmp += delta[offset] * key[offsetKey];
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offsetQuery = depthPrevQuery +
+        nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem;
+    
+    if (dirty)
+    {
+        query[offsetQuery] = tmp;
+    }
+    else
+    {
+        query[offsetQuery] += tmp;
+    }
+}
+
+kernel void queryCausalQuerySeq4BackwardFloat(
+    const device float * delta,
+    const device float4 * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device float4 * query,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / (size / 4);
+    uint j = id[0] % (size / 4);
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || j * 4 >= size ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    uint depthPrevKey = j * 4 + headKey * size;
+    uint depthPrevQuery = j * 4 + headQuery * size;
+    
+    float4 tmp = 0.0;
+    for (uint seqK=0; seqK<=seqQ; seqK++)
+    {
+        uint offset = seqK + headQuery * sequence +
+            nbNeurons * seqQ + sequence * nbNeurons * elem;
+        uint offsetKey = (depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+        
+        tmp += delta[offset] * key[offsetKey];
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offsetQuery = (depthPrevQuery +
+        nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem) / 4;
+    
+    if (dirty)
+    {
+        query[offsetQuery] = tmp;
+    }
+    else
+    {
+        query[offsetQuery] += tmp;
+    }
+}
+
+kernel void queryCausalKeySeqBackwardFloat(
+    const device float * delta,
+    const device float * query,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device float * key,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevKey / nbHeadsKey;
+    
+    uint headKey = id[0] / size;
+    uint j = id[0] % size;
+    uint elem = id[1] / sequence;
+    uint seqK = id[1] % sequence;
+    
+    if (headKey >= nbHeadsKey || j >= size ||
+        elem >= nbBatch || seqK >= sequence)
+    {
+        return ;
+    }
+    
+    uint nbBlocksHead = nbHeadsQuery == nbHeadsKey ?
+        1 : nbHeadsQuery / nbHeadsKey;
+    uint depthPrevKey = j + headKey * size;
+    
+    float tmp = 0.0;
+    for (uint blockHead=0; blockHead<nbBlocksHead; blockHead++)
+    {
+        uint headQuery = blockHead + nbBlocksHead * headKey;
+        uint depthPrevQuery = j + headQuery * size;
+        
+        for (uint seqQ=seqK; seqQ<sequence; seqQ++)
+        {
+            uint offset = seqK + headQuery * sequence +
+                nbNeurons * seqQ + sequence * nbNeurons * elem;
+            uint offsetQuery = depthPrevQuery +
+                nbNeuronsPrevQuery * seqQ +
+                sequence * nbNeuronsPrevQuery * elem;
+            
+            tmp += delta[offset] * query[offsetQuery];
+        }
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offsetKey = depthPrevKey +
+        nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+    
+    if (dirty)
+    {
+        key[offsetKey] = tmp;
+    }
+    else
+    {
+        key[offsetKey] += tmp;
+    }
+}
+
+kernel void queryCausalKeySeq4BackwardFloat(
+    const device float * delta,
+    const device float4 * query,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device float4 * key,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevKey / nbHeadsKey;
+    
+    uint headKey = id[0] / (size / 4);
+    uint j = id[0] % (size / 4);
+    uint elem = id[1] / sequence;
+    uint seqK = id[1] % sequence;
+    
+    if (headKey >= nbHeadsKey || j * 4 >= size ||
+        elem >= nbBatch || seqK >= sequence)
+    {
+        return ;
+    }
+    
+    uint nbBlocksHead = nbHeadsQuery == nbHeadsKey ?
+        1 : nbHeadsQuery / nbHeadsKey;
+    uint depthPrevKey = j * 4 + headKey * size;
+    
+    float4 tmp = 0.0;
+    for (uint blockHead=0; blockHead<nbBlocksHead; blockHead++)
+    {
+        uint headQuery = blockHead + nbBlocksHead * headKey;
+        uint depthPrevQuery = j * 4 + headQuery * size;
+        
+        for (uint seqQ=seqK; seqQ<sequence; seqQ++)
+        {
+            uint offset = seqK + headQuery * sequence +
+                nbNeurons * seqQ + sequence * nbNeurons * elem;
+            uint offsetQuery = (depthPrevQuery +
+                nbNeuronsPrevQuery * seqQ +
+                sequence * nbNeuronsPrevQuery * elem) / 4;
+            
+            tmp += delta[offset] * query[offsetQuery];
+        }
+    }
+    tmp /= sqrt((float)size);
+    
+    uint offsetKey = (depthPrevKey +
+        nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+    
+    if (dirty)
+    {
+        key[offsetKey] = tmp;
+    }
+    else
+    {
+        key[offsetKey] += tmp;
+    }
+}
diff --git a/Sources/GrAIdient/Metal/Kernel/NLPHalf.metal b/Sources/GrAIdient/Metal/Kernel/NLPHalf.metal
new file mode 100644
index 00000000..fea368fa
--- /dev/null
+++ b/Sources/GrAIdient/Metal/Kernel/NLPHalf.metal
@@ -0,0 +1,497 @@
+//
+// NLPHalf.metal
+// GrAIdient
+//
+// Created by Jean-François Reboud on 25/06/2024.
+//
+
+#include <metal_stdlib>
+using namespace metal;
+
+kernel void createRoPESeqMatrixHalf(
+    constant int * seqPositions,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & sequence,
+    device half * rotationMatrix,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint block = id[0];
+    uint seq = id[1];
+    
+    if (block >= nbBlocks || seq >= sequence)
+    {
+        return ;
+    }
+    
+    float position = (float)seqPositions[seq];
+    float theta = pow(
+        10000.0,
+        -2.0 * (float)block / (float)size
+    );
+    float mTheta = position * theta;
+    float cosVal = cos(mTheta);
+    float sinVal = sin(mTheta);
+    
+    uint offset = 2 * block + seq * size;
+    rotationMatrix[offset] = cosVal;
+    rotationMatrix[1 + offset] = sinVal;
+}
+
+kernel void RoPESeqForwardHalf(
+    const device half * outsPrev,
+    const device half * rotationMatrix,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device half * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint head = id[0] / nbBlocks;
+    uint block = id[0] % nbBlocks;
+    uint elem = id[1] / sequence;
+    uint seq = id[1] % sequence;
+    
+    if (head >= nbHeads || block >= nbBlocks ||
+        elem >= nbBatch || seq >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset1 = 2 * block + seq * size;
+    uint offset2 = 2 * block + head * size +
+        nbNeurons * seq + sequence * nbNeurons * elem;
+    
+    half cosVal = rotationMatrix[offset1];
+    half sinVal = rotationMatrix[1 + offset1];
+    
+    half in1 = outsPrev[offset2];
+    half in2 = outsPrev[1 + offset2];
+    
+    half out1 = in1 * cosVal - in2 * sinVal;
+    half out2 = in1 * sinVal + in2 * cosVal;
+    
+    outs[offset2] = out1;
+    outs[1 + offset2] = out2;
+}
+
+kernel void RoPESeqSeqBackwardHalf(
+    const device half * delta,
+    const device half * rotationMatrix,
+    constant uint & nbHeads,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device half * deltaPrev,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeurons / nbHeads;
+    uint nbBlocks = size / 2;
+    
+    uint head = id[0] / nbBlocks;
+    uint block = id[0] % nbBlocks;
+    uint elem = id[1] / sequence;
+    uint seq = id[1] % sequence;
+    
+    if (head >= nbHeads || block >= nbBlocks ||
+        elem >= nbBatch || seq >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset1 = 2 * block + seq * size;
+    uint offset2 = 2 * block + head * size +
+        nbNeurons * seq + sequence * nbNeurons * elem;
+    
+    half cosVal = rotationMatrix[offset1];
+    half sinVal = rotationMatrix[1 + offset1];
+    
+    half out1 = delta[offset2];
+    half out2 = delta[1 + offset2];
+    
+    half in1 = out1 * cosVal + out2 * sinVal;
+    half in2 = -out1 * sinVal + out2 * cosVal;
+    
+    if (dirty)
+    {
+        deltaPrev[offset2] = in1;
+        deltaPrev[1 + offset2] = in2;
+    }
+    else
+    {
+        deltaPrev[offset2] += in1;
+        deltaPrev[1 + offset2] += in2;
+    }
+}
+
+kernel void encodeCausalityHalf(
+    constant uint & nbHeadsQuery,
+    constant uint & nbNeurons,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device half * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery ||
+        seqK >= sequence || seqK <= seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = -1e4;
+}
+
+kernel void queryCausalSeqForwardHalf(
+    const device half * query,
+    const device half * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device half * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || seqK > seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    half tmp = 0.0;
+    
+    for (uint j=0; j<size; j++)
+    {
+        uint depthPrevKey = j + headKey * size;
+        uint depthPrevQuery = j + headQuery * size;
+        
+        uint offsetQuery = depthPrevQuery +
+            nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem;
+        uint offsetKey = depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+        
+        tmp += query[offsetQuery] * key[offsetKey];
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = tmp;
+}
+
+kernel void queryCausalSeq4ForwardHalf(
+    const device half4 * query,
+    const device half4 * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    device half * outs,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / sequence;
+    uint seqK = id[0] % sequence;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || seqK > seqQ ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    half4 tmp = 0.0;
+    
+    for (uint j=0; j<size/4; j++)
+    {
+        uint depthPrevKey = j * 4 + headKey * size;
+        uint depthPrevQuery = j * 4 + headQuery * size;
+        
+        uint offsetQuery = (depthPrevQuery +
+            nbNeuronsPrevQuery * seqQ +
+            sequence * nbNeuronsPrevQuery * elem) / 4;
+        uint offsetKey = (depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+        
+        tmp += query[offsetQuery] * key[offsetKey];
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offset = seqK + headQuery * sequence +
+        nbNeurons * seqQ + sequence * nbNeurons * elem;
+    outs[offset] = tmp[0] + tmp[1] + tmp[2] + tmp[3];
+}
+
+kernel void queryCausalQuerySeqBackwardHalf(
+    const device half * delta,
+    const device half * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device half * query,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / size;
+    uint j = id[0] % size;
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || j >= size ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    uint depthPrevKey = j + headKey * size;
+    uint depthPrevQuery = j + headQuery * size;
+    
+    half tmp = 0.0;
+    for (uint seqK=0; seqK<=seqQ; seqK++)
+    {
+        uint offset = seqK + headQuery * sequence +
+            nbNeurons * seqQ + sequence * nbNeurons * elem;
+        uint offsetKey = depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+        
+        tmp += delta[offset] * key[offsetKey];
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offsetQuery = depthPrevQuery +
+        nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem;
+    
+    if (dirty)
+    {
+        query[offsetQuery] = tmp;
+    }
+    else
+    {
+        query[offsetQuery] += tmp;
+    }
+}
+
+kernel void queryCausalQuerySeq4BackwardHalf(
+    const device half * delta,
+    const device half4 * key,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device half4 * query,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevQuery / nbHeadsQuery;
+    
+    uint headQuery = id[0] / (size / 4);
+    uint j = id[0] % (size / 4);
+    uint elem = id[1] / sequence;
+    uint seqQ = id[1] % sequence;
+    
+    if (headQuery >= nbHeadsQuery || j * 4 >= size ||
+        elem >= nbBatch || seqQ >= sequence)
+    {
+        return ;
+    }
+    
+    uint headKey = nbHeadsQuery == nbHeadsKey ?
+        headQuery : headQuery / nbHeadsKey;
+    uint depthPrevKey = j * 4 + headKey * size;
+    uint depthPrevQuery = j * 4 + headQuery * size;
+    
+    half4 tmp = 0.0;
+    for (uint seqK=0; seqK<=seqQ; seqK++)
+    {
+        uint offset = seqK + headQuery * sequence +
+            nbNeurons * seqQ + sequence * nbNeurons * elem;
+        uint offsetKey = (depthPrevKey +
+            nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+        
+        tmp += delta[offset] * key[offsetKey];
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offsetQuery = (depthPrevQuery +
+        nbNeuronsPrevQuery * seqQ + sequence * nbNeuronsPrevQuery * elem) / 4;
+    
+    if (dirty)
+    {
+        query[offsetQuery] = tmp;
+    }
+    else
+    {
+        query[offsetQuery] += tmp;
+    }
+}
+
+kernel void queryCausalKeySeqBackwardHalf(
+    const device half * delta,
+    const device half * query,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device half * key,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevKey / nbHeadsKey;
+    
+    uint headKey = id[0] / size;
+    uint j = id[0] % size;
+    uint elem = id[1] / sequence;
+    uint seqK = id[1] % sequence;
+    
+    if (headKey >= nbHeadsKey || j >= size ||
+        elem >= nbBatch || seqK >= sequence)
+    {
+        return ;
+    }
+    
+    uint nbBlocksHead = nbHeadsQuery == nbHeadsKey ?
+        1 : nbHeadsQuery / nbHeadsKey;
+    uint depthPrevKey = j + headKey * size;
+    
+    half tmp = 0.0;
+    for (uint blockHead=0; blockHead<nbBlocksHead; blockHead++)
+    {
+        uint headQuery = blockHead + nbBlocksHead * headKey;
+        uint depthPrevQuery = j + headQuery * size;
+        
+        for (uint seqQ=seqK; seqQ<sequence; seqQ++)
+        {
+            uint offset = seqK + headQuery * sequence +
+                nbNeurons * seqQ + sequence * nbNeurons * elem;
+            uint offsetQuery = depthPrevQuery +
+                nbNeuronsPrevQuery * seqQ +
+                sequence * nbNeuronsPrevQuery * elem;
+            
+            tmp += delta[offset] * query[offsetQuery];
+        }
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offsetKey = depthPrevKey +
+        nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem;
+    
+    if (dirty)
+    {
+        key[offsetKey] = tmp;
+    }
+    else
+    {
+        key[offsetKey] += tmp;
+    }
+}
+
+kernel void queryCausalKeySeq4BackwardHalf(
+    const device half * delta,
+    const device half4 * query,
+    constant uint & nbHeadsQuery,
+    constant uint & nbHeadsKey,
+    constant uint & nbNeurons,
+    constant uint & nbNeuronsPrevQuery,
+    constant uint & nbNeuronsPrevKey,
+    constant uint & nbBatch,
+    constant uint & sequence,
+    constant uint & dirty,
+    device half4 * key,
+    uint2 id [[ thread_position_in_grid ]])
+{
+    uint size = nbNeuronsPrevKey / nbHeadsKey;
+    
+    uint headKey = id[0] / (size / 4);
+    uint j = id[0] % (size / 4);
+    uint elem = id[1] / sequence;
+    uint seqK = id[1] % sequence;
+    
+    if (headKey >= nbHeadsKey || j * 4 >= size ||
+        elem >= nbBatch || seqK >= sequence)
+    {
+        return ;
+    }
+    
+    uint nbBlocksHead = nbHeadsQuery == nbHeadsKey ?
+        1 : nbHeadsQuery / nbHeadsKey;
+    uint depthPrevKey = j * 4 + headKey * size;
+    
+    half4 tmp = 0.0;
+    for (uint blockHead=0; blockHead<nbBlocksHead; blockHead++)
+    {
+        uint headQuery = blockHead + nbBlocksHead * headKey;
+        uint depthPrevQuery = j * 4 + headQuery * size;
+        
+        for (uint seqQ=seqK; seqQ<sequence; seqQ++)
+        {
+            uint offset = seqK + headQuery * sequence +
+                nbNeurons * seqQ + sequence * nbNeurons * elem;
+            uint offsetQuery = (depthPrevQuery +
+                nbNeuronsPrevQuery * seqQ +
+                sequence * nbNeuronsPrevQuery * elem) / 4;
+            
+            tmp += delta[offset] * query[offsetQuery];
+        }
+    }
+    tmp /= sqrt((half)size);
+    
+    uint offsetKey = (depthPrevKey +
+        nbNeuronsPrevKey * seqK + sequence * nbNeuronsPrevKey * elem) / 4;
+    
+    if (dirty)
+    {
+        key[offsetKey] = tmp;
+    }
+    else
+    {
+        key[offsetKey] += tmp;
+    }
+}
diff --git a/Sources/GrAIdient/Metal/MetalConfig.swift b/Sources/GrAIdient/Metal/MetalConfig.swift
index 76c91bde..139a0567 100644
--- a/Sources/GrAIdient/Metal/MetalConfig.swift
+++ b/Sources/GrAIdient/Metal/MetalConfig.swift
@@ -433,9 +433,6 @@ let CONFIG_KERNELS =
         "selectSeqForwardFloat",
         "selectSeqBackwardFloat",
         "layerCAMSeqForwardFloat",
-        "createRoPESeqMatrixFloat",
-        "RoPESeqForwardFloat",
-        "RoPESeqSeqBackwardFloat",
     ],
     "LayerSeqHalf": [
         "avgPoolSeqForwardHalf",
@@ -483,9 +480,30 @@ let CONFIG_KERNELS =
         "selectSeqForwardHalf",
         "selectSeqBackwardHalf",
         "layerCAMSeqForwardHalf",
+    ],
+    "NLPFloat": [
+        "createRoPESeqMatrixFloat",
+        "RoPESeqForwardFloat",
+        "RoPESeqSeqBackwardFloat",
+        "encodeCausalityFloat",
+        "queryCausalSeqForwardFloat",
+        "queryCausalSeq4ForwardFloat",
+        "queryCausalQuerySeqBackwardFloat",
+        "queryCausalQuerySeq4BackwardFloat",
+        "queryCausalKeySeqBackwardFloat",
+        "queryCausalKeySeq4BackwardFloat",
+    ],
+    "NLPHalf": [
         "createRoPESeqMatrixHalf",
         "RoPESeqForwardHalf",
         "RoPESeqSeqBackwardHalf",
+        "encodeCausalityHalf",
+        "queryCausalSeqForwardHalf",
+        "queryCausalSeq4ForwardHalf",
+        "queryCausalQuerySeqBackwardHalf",
+        "queryCausalQuerySeq4BackwardHalf",
+        "queryCausalKeySeqBackwardHalf",
+        "queryCausalKeySeq4BackwardHalf",
     ],
     "OptimizerFloat": [
         "clipGradientsFloat",
diff --git a/Sources/GrAIdient/Utils/Serialization.swift b/Sources/GrAIdient/Utils/Serialization.swift
index 25965f1f..ce056948 100644
--- a/Sources/GrAIdient/Utils/Serialization.swift
+++ b/Sources/GrAIdient/Utils/Serialization.swift
@@ -77,6 +77,7 @@ let LAYER_REGISTRY: [String: Codable.Type] = buildRegistry(
     MSE2D.self,
     Multiply2D.self,
     Pad2D.self,
+    QueryCausalSeq.self,
     QuerySeq.self,
     QuerySelfSeq.self,
     ResizeBilinear.self,
diff --git a/Tests/GrAIExamples/Base/python_lib/nlp/model.py b/Tests/GrAIExamples/Base/python_lib/nlp/model.py
index 9cb4e414..e7cd4920 100644
--- a/Tests/GrAIExamples/Base/python_lib/nlp/model.py
+++ b/Tests/GrAIExamples/Base/python_lib/nlp/model.py
@@ -240,18 +240,18 @@ def repeat(a):
             queries = torch.einsum("bhlj,lij->bhli", [queries, rotation_matrix])
             keys = torch.einsum("bhlj,lij->bhli", [keys, rotation_matrix])
 
-        """scores = torch.matmul(queries, keys.transpose(2, 3)) * self.scale
+        scores = torch.matmul(queries, keys.transpose(2, 3)) * self.scale
         if mask is not None:
             scores += mask
         scores = torch.softmax(
             scores.type(torch.float32), dim=-1
         ).type_as(scores)
 
-        output = torch.matmul(scores, values)
+        """output = torch.matmul(scores, values)
         output = output.transpose(1, 2).contiguous().reshape(B, L, -1)
 
         return self.wo(output), (keys, values)"""
-        return queries.transpose(1, 2).contiguous().reshape(B, L, -1), (keys, values)
+        return scores.transpose(1, 2).contiguous().reshape(B, L, -1), (keys, values)
 
 
 class FeedForward(torch.nn.Module):
diff --git a/Tests/GrAIExamples/NLPExample.swift b/Tests/GrAIExamples/NLPExample.swift
index 8e24a925..26decf00 100644
--- a/Tests/GrAIExamples/NLPExample.swift
+++ b/Tests/GrAIExamples/NLPExample.swift
@@ -16,7 +16,7 @@ final class NLPExample: XCTestCase
     let _modelPath = "TO/UPDATE"
     
     /// Prompt.
-    let _prompt = "I"
+    let _prompt = "Hello"
     
     /// Initialize test.
     override func setUp()
@@ -34,7 +34,9 @@ final class NLPExample: XCTestCase
     /// - Parameters:
     ///     - sequence: Length of the sequence.
     ///     - hiddenDim: Dimension of neurons in the main branch.
-    ///     - nbHeads:  Number of heads (groups) of neurons.
+    ///     - headDim: Dimension of neurons in the transformer branches.
+    ///     - nbHeads:  Number of heads (groups) of neurons for queries.
+    ///     - nbHeadsKV: Number of heads (groups) of neurons for keys and values.
     ///     - vocabularySize: Vocabulary size.
     /// - Returns: The model built.
     ///
@@ -42,7 +44,9 @@ final class NLPExample: XCTestCase
         modelPath: String,
         sequence: Int,
         hiddenDim: Int,
-        nbHeads: Int,
+        headDim: Int,
+        nbHeadsQuery: Int,
+        nbHeadsKV: Int,
         vocabularySize: Int) -> Model
     {
         let context = ModelContext(name: "NLP", curID: 0)
@@ -54,18 +58,42 @@ final class NLPExample: XCTestCase
             nbNeurons: hiddenDim, params: params
         )
         
-        layer = FullyConnectedSeq(
+        var query: LayerSeq = FullyConnectedSeq(
             layerPrev: layer,
-            nbNeurons: hiddenDim,
+            nbNeurons: nbHeadsQuery * headDim,
             activation: nil,
             biases: false,
             params: params
         )
+        query = try! RoPESeq(
+            layerPrev: query,
+            seqPositions: [Int](1...sequence),
+            nbHeads: nbHeadsQuery,
+            params: params
+        )
         
-        layer = try! RoPESeq(
+        var key: LayerSeq = FullyConnectedSeq(
             layerPrev: layer,
+            nbNeurons: nbHeadsKV * headDim,
+            activation: nil,
+            biases: false,
+            params: params
+        )
+        key = try! RoPESeq(
+            layerPrev: key,
             seqPositions: [Int](1...sequence),
-            nbHeads: nbHeads,
+            nbHeads: nbHeadsKV,
+            params: params
+        )
+        
+        layer = try! QueryCausalSeq(
+            query: query, key: key,
+            nbHeadsQuery: nbHeadsQuery, nbHeadsKey: nbHeadsKV, 
+            params: params
+        )
+        layer = try! SoftmaxSeq(
+            layerPrev: layer,
+            nbHeads: nbHeadsQuery,
             params: params
         )
         
@@ -146,7 +174,9 @@ final class NLPExample: XCTestCase
             modelPath: _modelPath,
             sequence: prompt.count,
             hiddenDim: 4096,
-            nbHeads: 32,
+            headDim: 128,
+            nbHeadsQuery: 32,
+            nbHeadsKV: 8,
             vocabularySize: 32000
         )
         
@@ -167,8 +197,15 @@ final class NLPExample: XCTestCase
         // Compare difference.
         for (elemOut, elemRef) in zip(arrayOut, arrayRef)
         {
-            let diffPercent = abs(elemOut - elemRef) / elemRef * 100.0
-            XCTAssert(diffPercent < 1)
+            if elemRef == 0.0
+            {
+                XCTAssert(elemOut == 0.0)
+            }
+            else
+            {
+                let diffPercent = abs(elemOut - elemRef) / elemRef * 100.0
+                XCTAssert(diffPercent < 1)
+            }
         }
     }
 }
diff --git a/Tests/GrAITests/Base/InputSeq/EmbeddingSeqMSE1DCase.swift b/Tests/GrAITests/Base/InputSeq/EmbeddingSeqMSE1DCase.swift
index 3a349b17..e1d62089 100644
--- a/Tests/GrAITests/Base/InputSeq/EmbeddingSeqMSE1DCase.swift
+++ b/Tests/GrAITests/Base/InputSeq/EmbeddingSeqMSE1DCase.swift
@@ -28,8 +28,8 @@ class EmbeddingSeqMSE1DCase: XCTestCase, Input1DCase, IOCase
     override func setUp()
     {
         batchSize = 5
-        sequence = 7
-        vocabularySize = 120
+        sequence = 5
+        vocabularySize = 7
         _ = MetalKernel.get
         
         GrAI.Opti.GPU = true
diff --git a/Tests/GrAITests/NLPTests.swift b/Tests/GrAITests/NLPTests.swift
index 0ad3ca97..01372740 100644
--- a/Tests/GrAITests/NLPTests.swift
+++ b/Tests/GrAITests/NLPTests.swift
@@ -73,6 +73,58 @@ class NLPGradTests: EmbeddingSeqMSE1DCase
                 nbHeads: 3,
                 params: params
             )
+        
+        case "QueryCausal1":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 3, nbHeadsKey: 3,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 3,
+                params: params
+            )
+            
+        case "QueryCausal2":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 2 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 4 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 4, nbHeadsKey: 2,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 4,
+                params: params
+            )
             
         default:
             fatalError("Unreachable.")
@@ -133,6 +185,32 @@ class NLPGradTests: EmbeddingSeqMSE1DCase
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    func testQueryCausal1CPU() throws
+    {
+        GrAI.Opti.CPU = true
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    func testQueryCausal1GPU() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    func testQueryCausal2CPU() throws
+    {
+        GrAI.Opti.CPU = true
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
+    
+    func testQueryCausal2GPU() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -192,6 +270,58 @@ class NLPFlowTests: EmbeddingSeqMSE1DCase
                 params: params
             )
             
+        case "QueryCausal1":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 3, nbHeadsKey: 3,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 3,
+                params: params
+            )
+            
+        case "QueryCausal2":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 2 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 4 * 3,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 4, nbHeadsKey: 2,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 4,
+                params: params
+            )
+            
         default:
             fatalError("Unreachable.")
         }
@@ -230,6 +360,18 @@ class NLPFlowTests: EmbeddingSeqMSE1DCase
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -276,6 +418,162 @@ class NLPFlowPrecisionTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer, diffThreshold: 0.002)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer, diffThreshold: 0.002)
+    }
+}
+
+// -----------------------------------------------------------------------------
+// Compare GPU gradients with CPU ones through time.
+// We expect to see errors ~ 1e-7 and less.
+// -----------------------------------------------------------------------------
+class NLP4FlowTests: EmbeddingSeqMSE1DCase
+{
+    private func _buildTrainer(_ model: String) -> FlowTrainer
+    {
+        let trainer = FlowTrainer(
+            name: "NLP",
+            params: optimizerParams
+        )
+        trainer.build()
+        {
+            (context: ModelContext) in
+            buildModel(model: model, context: context)
+        }
+        return trainer
+    }
+    
+    func buildModel(model: String, context: ModelContext)
+    {
+        let params = GrAI.Model.Params(context: context)
+        
+        var layer: LayerSeq = EmbeddingSeq(
+            sequence: sequence,
+            vocabularySize: vocabularySize,
+            nbNeurons: 4, params: params
+        )
+        
+        switch model
+        {
+        case "QueryCausal1":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 4,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 3 * 4,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 3, nbHeadsKey: 3,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 3,
+                params: params
+            )
+            
+        case "QueryCausal2":
+            let otherLayer: LayerSeq = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 2 * 4,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = FullyConnectedSeq(
+                layerPrev: layer,
+                nbNeurons: 4 * 4,
+                activation: nil,
+                biases: false,
+                params: params
+            )
+            layer = try! QueryCausalSeq(
+                query: layer, key: otherLayer,
+                nbHeadsQuery: 4, nbHeadsKey: 2,
+                params: params
+            )
+            layer = try! SoftmaxSeq(
+                layerPrev: layer,
+                nbHeads: 4,
+                params: params
+            )
+            
+        default:
+            fatalError("Unreachable.")
+        }
+        
+        var head: Layer1D = AvgPoolSeq(layerPrev: layer, params: params)
+        
+        head = try! FullyConnected(
+            layerPrev: head, nbNeurons: 1,
+            activation: LeakyReLU.str, biases: true, params: params
+        )
+        
+        _ = MSE1D(layerPrev: head, params: params)
+    }
+    
+    func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
+}
+
+// -----------------------------------------------------------------------------
+// Compare GPU gradients with Float precision versus Float16 precision.
+// We expect to see errors ~ 1e-4 and less.
+// -----------------------------------------------------------------------------
+class NLP4FlowPrecisionTests: NLP4FlowTests
+{
+    private func _buildTrainer(_ model: String) -> FlowPrecisionTrainer
+    {
+        let trainer = FlowPrecisionTrainer(
+            name: "NLP",
+            params: optimizerParams
+        )
+        trainer.build()
+        {
+            (context: ModelContext) in
+            buildModel(model: model, context: context)
+        }
+        return trainer
+    }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer, diffThreshold: 0.002)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer, diffThreshold: 0.002)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -330,6 +628,18 @@ class NLPFlowResetTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -384,6 +694,18 @@ class NLPFlowReverseTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -506,6 +828,18 @@ class NLPInferenceTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -553,6 +887,18 @@ class NLPLoadTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -644,4 +990,16 @@ class NLPTransformTests: NLPFlowTests
         let trainer = _buildTrainer("RoPE")
         run(trainer)
     }
+    
+    override func testQueryCausal1() throws
+    {
+        let trainer = _buildTrainer("QueryCausal1")
+        run(trainer)
+    }
+    
+    override func testQueryCausal2() throws
+    {
+        let trainer = _buildTrainer("QueryCausal2")
+        run(trainer)
+    }
 }