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
102 lines
3.6 KiB
Swift
102 lines
3.6 KiB
Swift
import Foundation
|
|
import MarkBase
|
|
import Metal
|
|
import RDMAKit
|
|
|
|
// ─────────────────────────────────────────────────────────────
|
|
// RDMA Distribution Service — Pipeline Parallelism
|
|
// ─────────────────────────────────────────────────────────────
|
|
|
|
public actor RDMADistributionService {
|
|
public enum Role: Sendable {
|
|
case primary(splitLayer: Int)
|
|
case secondary
|
|
}
|
|
|
|
public struct Config: Sendable {
|
|
public let role: Role
|
|
public let peerAddress: String?
|
|
public let peerPort: UInt16
|
|
|
|
public init(role: Role, peerAddress: String? = nil, peerPort: UInt16 = 0) {
|
|
self.role = role
|
|
self.peerAddress = peerAddress
|
|
self.peerPort = peerPort
|
|
}
|
|
}
|
|
|
|
// RDMA state
|
|
private var context: RDMAContext?
|
|
private var pd: RDMAProtectionDomain?
|
|
private var sendCQ: RDMACompletionQueue?
|
|
private var recvCQ: RDMACompletionQueue?
|
|
private var qp: RDMAQueuePair?
|
|
|
|
// Registered GPU buffers
|
|
private var registeredHiddenState: RegisteredMemory?
|
|
|
|
public let config: Config
|
|
public let discovery: RDMADiscovery
|
|
|
|
public init(config: Config) {
|
|
self.config = config
|
|
self.discovery = RDMADiscovery()
|
|
}
|
|
|
|
/// Initialize RDMA device and resources. Returns false if no RDMA device is available.
|
|
@discardableResult
|
|
public func initialize() throws -> Bool {
|
|
let devices = discovery.listDevices()
|
|
guard let device = devices.first else {
|
|
print(" RDMA: No devices found — running in local-only mode")
|
|
return false
|
|
}
|
|
|
|
let tbInfo = device.isThunderbolt ? " (Thunderbolt)" : ""
|
|
print(" RDMA: Found device '\(device.name)'\(tbInfo)")
|
|
context = try discovery.openDevice(device)
|
|
let attrs = try context!.queryDeviceAttributes()
|
|
print(" RDMA: Firmware \(attrs.fwVer), max MR size \(attrs.maxMRSize) bytes")
|
|
|
|
pd = try RDMAProtectionDomain(context: context!)
|
|
sendCQ = try RDMACompletionQueue(context: context!, cqe: 128)
|
|
recvCQ = try RDMACompletionQueue(context: context!, cqe: 128)
|
|
qp = try RDMAQueuePair(pd: pd!, sendCQ: sendCQ!, recvCQ: recvCQ!)
|
|
try qp!.modifyToInit()
|
|
print(" RDMA: QP initialized in RC mode")
|
|
|
|
if case .primary = config.role {
|
|
try qp!.modifyToRTR(destQPN: 0)
|
|
try qp!.modifyToRTS()
|
|
}
|
|
|
|
print(" RDMA: Ready")
|
|
return true
|
|
}
|
|
|
|
/// Register a Metal buffer for remote RDMA access
|
|
public func registerBuffer(_ buffer: MTLBuffer) throws -> RegisteredMemory? {
|
|
guard let pd else { return nil }
|
|
let reg = try pd.registerMTLBuffer(
|
|
UnsafeMutableRawPointer(buffer.contents()),
|
|
length: buffer.length
|
|
)
|
|
print(" RDMA: Registered buffer \(buffer.length) bytes (lkey=\(reg.lkey), rkey=\(reg.rkey))")
|
|
return reg
|
|
}
|
|
|
|
/// Register the model's hidden state buffer for pipeline-split transfer
|
|
public func registerHiddenState(_ model: E4BModel) throws {
|
|
let buf = model.temps.io
|
|
let reg = try registerBuffer(buf)
|
|
registeredHiddenState = reg
|
|
print(" RDMA: Hidden state buffer registered (\(buf.length) bytes)")
|
|
}
|
|
|
|
/// Split point for pipeline parallelism
|
|
public var splitLayer: Int? {
|
|
if case .primary(let layer) = config.role { return layer }
|
|
return nil
|
|
}
|
|
}
|