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
This commit is contained in:
@@ -0,0 +1,101 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user