v2: Initial clean branch with unit tests + CI/CD pipeline

- Started from ac75faa (initial E4B-MarkBase integration) - Kept Sources/ (all engine code) + Package.swift + .gitignore - Removed all ad-hoc tests, documentation, scripts, Python files - Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest) - Added .gitea/workflows/ci.yaml (build + unit tests + lint) - Added Scripts/check_resources.sh (memory-aware test runner) - Added Tests/Manifest.json (resource requirements for all tests) - Focus: 4-bit quantized models only
2026-07-05 13:29:25 +08:00
commit 8a66b9086a
90 changed files with 22252 additions and 0 deletions
@@ -0,0 +1,67 @@
+import Metal
+
+/// Per-layer KV cache supporting sliding window (rotating) and full (growing) attention.
+public final class KVCache {
+    let isSliding: Bool
+    let maxLength: Int
+    let nKvHeads: Int
+    let headDim: Int
+    let buffer: MTLBuffer  // contiguous K then V: [2 * maxLength * nKvHeads * headDim]
+    private(set) var currentLength: Int = 0
+
+    init(device: MTLDevice, isSliding: Bool, maxContextLength: Int, nKvHeads: Int, headDim: Int) {
+        self.isSliding = isSliding
+        self.maxLength = isSliding ? 512 : maxContextLength
+        self.nKvHeads = nKvHeads
+        self.headDim = headDim
+
+        let perStep = nKvHeads * headDim * MemoryLayout<Float>.stride
+        let total = 2 * self.maxLength * perStep
+        self.buffer = device.makeBuffer(length: total, options: .storageModeShared)!
+    }
+
+    var effectiveLength: Int {
+        isSliding ? min(currentLength, 512) : currentLength
+    }
+
+    /// Key buffer start offset (in bytes, from buffer start)
+    var keyBaseOffset: Int { 0 }
+
+    /// Value buffer start offset (in bytes, from buffer start) — immediately after K
+    var valueBaseOffset: Int {
+        maxLength * nKvHeads * headDim * MemoryLayout<Float>.stride
+    }
+
+    /// Byte offset for a given logical position in the key region.
+    func keyOffset(for position: Int) -> Int {
+        let p = isSliding ? (position % maxLength) : position
+        return p * nKvHeads * headDim * MemoryLayout<Float>.stride
+    }
+
+    func valueOffset(for position: Int) -> Int {
+        valueBaseOffset + keyOffset(for: position)
+    }
+
+    /// Store K,V into cache at the given logical position.
+    func store(key: MTLBuffer, keySrcOffset: Int,
+               value: MTLBuffer, valueSrcOffset: Int,
+               position: Int,
+               commandBuffer: MTLCommandBuffer) {
+        let stepBytes = nKvHeads * headDim * MemoryLayout<Float>.stride
+        currentLength = max(currentLength, position + 1)
+
+        let blit = commandBuffer.makeBlitCommandEncoder()!
+        blit.copy(from: key, sourceOffset: keySrcOffset,
+                  to: buffer, destinationOffset: keyOffset(for: position),
+                  size: stepBytes)
+        blit.copy(from: value, sourceOffset: valueSrcOffset,
+                  to: buffer, destinationOffset: valueOffset(for: position),
+                  size: stepBytes)
+        blit.endEncoding()
+    }
+    
+    /// Reset cache for new sequence
+    func reset() {
+        currentLength = 0
+    }
+}