MarkBase Admin
ac75faa0cc
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
68 lines
2.4 KiB
Swift
68 lines
2.4 KiB
Swift
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
|
|
}
|
|
}
|