From 4fc487b3344780d92bfd6f678ab4598fbe9e5ad7 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Wed, 21 Oct 2026 20:13:53 +0400 Subject: [PATCH 01/11] Add concurrent synchronous safetensors loading across files (#462) * Add concurrent synchronous safetensors loading across files The safetensors loader is lazy and a single eval of a whole checkpoint visits tensors in dictionary order, serializing reads at effectively random file offsets. Add loadArrays(urls:) and loadArraysAndMetadata(urls:), which split every file into contiguous byte-balanced ranges in file-offset order and evaluate the ranges from concurrent work items, so reads are sequential and header parsing, I/O, and copies overlap. This moves the technique from ml-explore/mlx-swift-lm#575 down into mlx-swift, per review feedback there. Measured on an M4 Pro with NVMe at ~6.1 GB/s: an 18 GB checkpoint loads 20% faster cold (4.5 to 5.6 GB/s) and a 10 GB checkpoint 5-8% faster warm; parity when the serial loader is already disk-bound. --- Source/MLX/IO+Concurrent.swift | 223 ++++++++++++++++ Source/MLX/Transforms+Compile.swift | 132 ++++----- Tests/MLXTests/ConcurrentLoadTests.swift | 266 +++++++++++++++++++ Tests/MLXTests/MaterializedHazardTests.swift | 149 ++++++++++- 4 files changed, 687 insertions(+), 83 deletions(-) create mode 100644 Source/MLX/IO+Concurrent.swift create mode 100644 Tests/MLXTests/ConcurrentLoadTests.swift diff --git a/Source/MLX/IO+Concurrent.swift b/Source/MLX/IO+Concurrent.swift new file mode 100644 index 00000000..e3b2a720 --- /dev/null +++ b/Source/MLX/IO+Concurrent.swift @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +import CoreFoundation +import Foundation + +// MARK: - Concurrent safetensors loading + +// Safetensors payload reads are lazy. Submit tensors in reverse file-offset order +// because MLX builds its evaluation tape in input order and consumes it from the back. +// Each Load primitive queues its read on MLX's I/O pool, which owns read concurrency; +// one checkedEval schedules the complete checkpoint and waits for all reads to finish. + +/// Match the limit enforced by MLX and the safetensors reference implementation. +private let maximumSafetensorHeaderBytes: UInt64 = 100_000_000 + +/// Decode a JSON integer without accepting booleans, floating-point values, or values +/// outside the supported signed 64-bit offset range. +private func safetensorOffset(_ value: Any) -> Int64? { + guard let number = value as? NSNumber else { return nil } + let numberType = String(cString: number.objCType) + guard + CFGetTypeID(number) != CFBooleanGetTypeID(), + numberType != "f", + numberType != "d", + numberType != "D" + else { + return nil + } + let offset = number.int64Value + return offset >= 0 ? offset : nil +} + +/// The tensor names in the safetensors file at `url`, ordered by position in the file. +/// +/// Reads the 8-byte header length and the JSON header only. Throws when the file is not a +/// well-formed safetensors file; callers fall back to loading the file whole. +func safetensorNamesInFileOrder(url: URL) throws -> [String] { + struct Malformed: Error {} + + let handle = try FileHandle(forReadingFrom: url) + defer { try? handle.close() } + + guard let lengthData = try handle.read(upToCount: 8), lengthData.count == 8 else { + throw Malformed() + } + let headerLength = lengthData.withUnsafeBytes { $0.loadUnaligned(as: UInt64.self) } + .littleEndian + guard headerLength > 0, headerLength < maximumSafetensorHeaderBytes else { + throw Malformed() + } + guard let headerData = try handle.read(upToCount: Int(headerLength)), + headerData.count == headerLength, + let header = try JSONSerialization.jsonObject(with: headerData) as? [String: Any] + else { + throw Malformed() + } + + var spans = [(name: String, begin: Int64, end: Int64)]() + for (name, value) in header { + guard name != "__metadata__" else { continue } + guard let entry = value as? [String: Any], + let offsets = entry["data_offsets"] as? [Any], offsets.count == 2, + let begin = safetensorOffset(offsets[0]), + let end = safetensorOffset(offsets[1]), + end >= begin + else { + throw Malformed() + } + spans.append((name, begin, end)) + } + // Safetensors requires the data buffer to be covered contiguously. Sorting by end as + // well handles zero-byte tensors that share the following tensor's begin offset. + spans.sort { + $0.begin == $1.begin ? $0.end < $1.end : $0.begin < $1.begin + } + var expectedBegin: Int64 = 0 + for span in spans { + guard span.begin == expectedBegin else { throw Malformed() } + expectedBegin = span.end + } + return spans.map(\.name) +} + +/// Roots in the order MLX must receive them to schedule reads in file-offset order. +func safetensorArraysInEvaluationOrder( + arrays: [[String: MLXArray]], names: [[String]?] +) -> [MLXArray] { + var orderedArrays = [MLXArray]() + for (fileArrays, fileNames) in zip(arrays, names) { + if let fileNames, + fileNames.count == fileArrays.count, + fileNames.allSatisfy({ fileArrays[$0] != nil }) + { + orderedArrays.append(contentsOf: fileNames.compactMap { fileArrays[$0] }) + } else { + // A planning failure or disagreement with MLX's header read must not leave + // any arrays out of the completion barrier. Their read order is unspecified. + orderedArrays.append(contentsOf: fileArrays.values) + } + } + // Reverse the whole checkpoint so MLX visits files in URL order and tensors within + // each file in ascending offset order. Keep duplicate names from different files. + return orderedArrays.reversed() +} + +/// Lock-guarded results of the parallel, header-only preparation pass. +private final class ConcurrentLoadPreparation: @unchecked Sendable { + private let lock = NSLock() + private var perFileArrays: [[String: MLXArray]] + private var perFileMetadata: [[String: String]] + private var namesPerFile: [[String]?] + private var firstError: Error? + + init(fileCount: Int) { + perFileArrays = Array(repeating: [:], count: fileCount) + perFileMetadata = Array(repeating: [:], count: fileCount) + namesPerFile = Array(repeating: nil, count: fileCount) + } + + func set( + file: Int, + arrays: [String: MLXArray], + metadata: [String: String], + names: [String]? + ) { + lock.lock() + defer { lock.unlock() } + perFileArrays[file] = arrays + perFileMetadata[file] = metadata + namesPerFile[file] = names + } + + func record(error: Error) { + lock.lock() + defer { lock.unlock() } + if firstError == nil { firstError = error } + } + + func result() throws -> ( + arrays: [[String: MLXArray]], + metadata: [[String: String]], + names: [[String]?] + ) { + lock.lock() + defer { lock.unlock() } + if let firstError { throw firstError } + return (perFileArrays, perFileMetadata, namesPerFile) + } +} + +/// Load a dictionary of ``MLXArray`` from several `safetensors` files at once, +/// scheduling payload reads in file-offset order and waiting for every array. +/// +/// The dictionaries are merged in file order: on a duplicate name the later file wins, +/// matching a serial load-and-merge loop over the same `urls`. Passing a single URL +/// also submits its tensors in file-offset order. MLX handles read concurrency. +/// +/// This function is synchronous and blocks until every file is loaded. From async code, +/// call it via a continuation on a `DispatchQueue` rather than directly on a Swift +/// concurrency cooperative thread. +/// +/// - Parameters: +/// - urls: URLs of the `safetensors` files to load +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``loadArrays(url:stream:)`` +/// - ``loadArraysAndMetadata(urls:stream:)`` +public func loadArrays(urls: [URL], stream: StreamOrDevice = .cpu) throws -> [String: MLXArray] { + try loadArraysAndMetadata(urls: urls, stream: stream).0 +} + +/// Load a dictionary of ``MLXArray`` and merged metadata `[String: String]` from several +/// `safetensors` files at once, scheduling payload reads in file-offset order. +/// +/// See ``loadArrays(urls:stream:)``. The metadata dictionaries are merged with the same +/// rule as the arrays: on a duplicate key the later file wins. +/// +/// - Parameters: +/// - urls: URLs of the `safetensors` files to load +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``loadArrays(urls:stream:)`` +/// - ``loadArraysAndMetadata(url:stream:)`` +public func loadArraysAndMetadata(urls: [URL], stream: StreamOrDevice = .cpu) throws -> ( + [String: MLXArray], [String: String] +) { + // Only header parsing happens on Dispatch worker threads. Payload reads stay lazy + // until checkedEval below, where MLX owns their concurrency. + let preparation = ConcurrentLoadPreparation(fileCount: urls.count) + DispatchQueue.concurrentPerform(iterations: urls.count) { file in + let url = urls[file] + let names = try? safetensorNamesInFileOrder(url: url) + do { + let (arrays, metadata) = try loadArraysAndMetadata(url: url, stream: stream) + preparation.set( + file: file, arrays: arrays, metadata: metadata, names: names) + } catch { + preparation.record(error: error) + } + } + + let prepared = try preparation.result() + let perFileArrays = prepared.arrays + let perFileMetadata = prepared.metadata + let evaluationOrder = safetensorArraysInEvaluationOrder( + arrays: perFileArrays, names: prepared.names) + + // checkedEval schedules the ordered roots under evalLock and waits outside it. + // Keep every root, including overwritten names, through this checked I/O barrier. + try checkedEval(evaluationOrder) + + var arrays = [String: MLXArray]() + var metadata = [String: String]() + for fileArrays in perFileArrays { + arrays.merge(fileArrays) { _, new in new } + } + for fileMetadata in perFileMetadata { + metadata.merge(fileMetadata) { _, new in new } + } + return (arrays, metadata) +} diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index d04944df..24ed86a8 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -3,17 +3,48 @@ import Cmlx import Foundation -// Note: this is all immutable state -- the `id` property is only set at init time -- -// except for ``CompiledFunction/tracedConstants``, which is guarded by -// ``CompiledFunction/lock``. -// Also: this is marked as Sendable even though it may capture non-Sendable state -// in inputs/outputs. The functions that create it will not return a Sendable -// type in that case -- see compile(inputs:) +/// A weak handle to MLX's thread-local cache. Retaining this wrapper does not keep +/// the backend cache alive after its thread exits. +private final class ThreadCompilationCache { + private let ctx: mlx_compile_cache + + private init() { + var cache = mlx_compile_cache_new() + mlx_detail_compile_cache(&cache) + self.ctx = cache + } + + deinit { + mlx_compile_cache_free(ctx) + } + + static var current: ThreadCompilationCache { + let key = "MLX.ThreadCompilationCache" + let dictionary = Thread.current.threadDictionary + if let cache = dictionary[key] as? ThreadCompilationCache { + return cache + } + let cache = ThreadCompilationCache() + dictionary[key] = cache + return cache + } + + func erase(_ functionID: UInt) { + mlx_detail_compile_erase(ctx, functionID) + } +} + +// Mutable cache bookkeeping is guarded by evalLock. The stateful overloads also +// hold the instance lock while replacing observed inputs and outputs. +// Only the overloads accepting Sendable captures return a Sendable closure. final class CompiledFunction: @unchecked (Sendable) { /// unique (for the lifetime of the object) identifier for the compiled function private var id: UInt! + /// Every backend cache this function has used, guarded by evalLock. + private var caches: [ObjectIdentifier: ThreadCompilationCache] = [:] + /// guards the observed state (``inputs`` / ``outputs``) -- see #226 /// /// This is always acquired *inside* `evalLock`; see ``call(_:)`` for the @@ -75,77 +106,24 @@ final class CompiledFunction: @unchecked (Sendable) { self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque()) } - /// Remove the compiled structure from the back end. - /// - /// ### Why this takes `evalLock` - /// - /// `mlx_detail_compile_erase` is `compiler_cache().erase(fun_id)`, i.e. - /// `std::unordered_map::erase` on the *same* process-global - /// `mlx::core::detail::CompilerCache` that ``innerCall(_:)`` reads and - /// inserts into via `CompilerCache::find` (`cache_[fun_id]`, which itself - /// can insert and rehash, followed by `entries.push_back`). That cache has - /// no synchronization of its own, and `find` hands back a *reference* to an - /// entry that the caller then fills in and reads from across the whole - /// trace. An unlocked erase can therefore rehash the map out from under a - /// concurrent `find`, or destroy the very entry another thread is compiling - /// into. - /// - /// Every other use of that cache happens under `evalLock`, so the erase has - /// to as well; otherwise the lock does not actually make the cache - /// single-threaded, it only makes most of it single-threaded. - /// - /// ### Why blocking here is safe - /// - /// `deinit` runs on whichever thread drops the last reference: - /// - /// - A thread that holds no MLX lock simply waits for the current eval or - /// compiled call to finish. Bounded, since `evalLock` is only ever held - /// for the duration of one call. - /// - A thread already inside a compiled call or an `eval` -- for instance a - /// traced body that drops the last reference to some other compiled - /// function -- already holds `evalLock`, and it is an `NSRecursiveLock`, - /// so re-entry is free. - /// - A thread holding some other ``CompiledFunction``'s ``lock`` also holds - /// `evalLock` already (see ``call(_:)``), so that is the same case; no new - /// edge is added to the lock graph. - /// - /// Within this library the only releases of a ``CompiledFunction`` are the - /// caller's own (the closure returned by ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` - /// holds the sole strong reference) and the `mlx_closure_free` of the - /// closure built in ``innerCall(_:)``, which runs on the calling thread with - /// `evalLock` already held. Nothing hands a reference to a back-end worker - /// thread, so the erase can never be the thing an `evalLock` holder is - /// waiting for. - /// - /// The erase stays synchronous rather than being deferred to a queue on - /// purpose: ``id`` is the object's own address, so once this object is - /// freed a new ``CompiledFunction`` can be allocated at the same address and - /// receive the same id. A deferred erase could delete that new function's - /// cache entries. + // Erase every calling thread's cache before this address can be reused as an ID. deinit { - let functionID = id! withEvalLock { #if DEBUG - // Trivially true two lines below the acquisition -- it is here to - // catch a future change that moves the erase back out from under - // the lock, the same way `withInstanceLock` guards `call(_:)`. EvalLockOwnership.requireHeldByCurrentThread( "the compiler cache was erased without holding evalLock; " + "see CompiledFunction.deinit") EvalLockOwnership.recordCompileErase() #endif - Self.eraseCache(functionID) + eraseCaches() } } - /// Remove the compiled structure for `functionID` from the back end. - /// - /// - Precondition: `evalLock` is held -- see ``deinit``. - private static func eraseCache(_ functionID: UInt) { - var cache = mlx_compile_cache_new() - mlx_detail_compile_cache(&cache) - defer { mlx_compile_cache_free(cache) } - mlx_detail_compile_erase(cache, functionID) + /// Must be called under evalLock, including when sealed state is rebound. + private func eraseCaches() { + for cache in caches.values { + cache.erase(id) + } } /// Evaluate the compiled function. @@ -155,10 +133,8 @@ final class CompiledFunction: @unchecked (Sendable) { /// `evalLock` is acquired **outermost**, before the per-instance ``lock``. /// That order is not a preference, it is forced: /// - /// - Tracing mutates process-global state that has no synchronization of its - /// own -- `mlx::core::detail::CompilerCache` is an unguarded singleton and - /// `detail::InTracing::trace_stack()` is a plain function-local `static`, - /// not `thread_local` -- so `evalLock` has to span the whole trace (#339). + /// - A trace can call back into Swift and mutate observed state, so + /// `evalLock` spans the whole trace and cache invalidation (#339). /// - The trace calls back into Swift, and the traced body routinely calls /// *other* compiled functions: every `MLXNN` activation is a file-scope /// `CompiledFunction`, so e.g. `silu` inside a traced body enters @@ -182,16 +158,7 @@ final class CompiledFunction: @unchecked (Sendable) { /// closure creation at the top of ``innerCall(_:)``. func call(_ arguments: [MLXArray]) -> [MLXArray] { withEvalLock { - // ``lock`` exists only to guard the observed state (#226): the - // `_updateInternal` swaps in `innerCall` are the sole mutation of - // anything reachable from `self` -- every other stored property is - // `let` except ``tracedConstants`` (only written when there is - // state), and `id` is written once at init. With no state there is - // nothing to guard, and `evalLock` above already serializes the - // compiler cache and every other caller. - // - // This is the common case: every compiled function in `MLXNN` is - // stateless. + // Stateless functions need only evalLock for cache bookkeeping. if inputs.isEmpty && outputs.isEmpty { return innerCall(arguments) } @@ -220,6 +187,9 @@ final class CompiledFunction: @unchecked (Sendable) { /// too unless this function has no observed state; see ``call(_:)``, /// which is the only caller. func innerCall(_ arguments: [MLXArray]) -> [MLXArray] { + let cache = ThreadCompilationCache.current + caches[ObjectIdentifier(cache)] = cache + // Sealed arrays (MaterializedArray, the parameters of a sealed module) // can never change, and must not be mutated, so they are not threaded // through the tracer swap below -- they are simply captured as @@ -239,7 +209,7 @@ final class CompiledFunction: @unchecked (Sendable) { traced.count == constants.count && zip(traced, constants).allSatisfy { $0 === $1 } if !same { - Self.eraseCache(id) + eraseCaches() tracedConstants = constants } } else { diff --git a/Tests/MLXTests/ConcurrentLoadTests.swift b/Tests/MLXTests/ConcurrentLoadTests.swift new file mode 100644 index 00000000..a9e27a7b --- /dev/null +++ b/Tests/MLXTests/ConcurrentLoadTests.swift @@ -0,0 +1,266 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import XCTest + +@testable import MLX + +final class ConcurrentLoadTests: XCTestCase { + + let temporaryPath = FileManager.default.temporaryDirectory.appending( + path: UUID().uuidString, + directoryHint: .isDirectory + ) + + override func setUpWithError() throws { + try FileManager.default.createDirectory( + at: temporaryPath, + withIntermediateDirectories: false + ) + } + + override func tearDownWithError() throws { + try FileManager.default.removeItem(at: temporaryPath) + } + + private func write( + arrays: [String: MLXArray], metadata: [String: String] = [:], name: String + ) throws -> URL { + let url = temporaryPath.appending(path: name, directoryHint: .notDirectory) + try save(arrays: arrays, metadata: metadata, url: url) + return url + } + + private func writeRawSafetensor( + header: String, payload: Data = Data(), name: String + ) throws -> URL { + let url = temporaryPath.appending(path: name, directoryHint: .notDirectory) + var headerLength = UInt64(header.utf8.count).littleEndian + var data = withUnsafeBytes(of: &headerLength) { Data($0) } + data.append(contentsOf: header.utf8) + data.append(payload) + try data.write(to: url) + return url + } + + func testLoadArraysMatchesSerialLoads() throws { + var urls = [URL]() + var expected = [String: MLXArray]() + for file in 0 ..< 3 { + var arrays = [String: MLXArray]() + // More roots than MLX's default BFS width exercises a larger evaluation tape. + for tensor in 0 ..< 32 { + let name = "file\(file).tensor\(tensor)" + arrays[name] = MLXArray(Int32(0) ..< 1024).reshaped(32, 32) + Int32(tensor) + } + urls.append(try write(arrays: arrays, name: "shard\(file).safetensors")) + expected.merge(arrays) { _, new in new } + } + + let loaded = try loadArrays(urls: urls) + + XCTAssertEqual(loaded.keys.sorted(), expected.keys.sorted()) + XCTAssertTrue(loaded.values.allSatisfy { isEvaluated($0) }) + for (name, array) in expected { + assertEqual(try XCTUnwrap(loaded[name]), array) + } + } + + func testLoadArraysSingleFile() throws { + let arrays = [ + "a": MLXArray(converting: [1.5, 2.5, 3.5]), + "b": MLX.ones([8, 8]), + ] + let url = try write(arrays: arrays, name: "single.safetensors") + + let loaded = try loadArrays(urls: [url]) + XCTAssertEqual(loaded.keys.sorted(), arrays.keys.sorted()) + for (name, array) in arrays { + assertEqual(try XCTUnwrap(loaded[name]), array) + } + } + + func testLoadArraysMaterializesBeforeReturning() throws { + let firstExpected = MLXArray(Int32(0) ..< 4096).reshaped(64, 64) + let secondExpected = MLXArray(Int32(4096) ..< 8192).reshaped(64, 64) + let firstURL = try write( + arrays: ["first": firstExpected], name: "first-materialized.safetensors") + let secondURL = try write( + arrays: ["second": secondExpected], name: "second-materialized.safetensors") + + let loaded = try loadArrays(urls: [firstURL, secondURL]) + for url in [firstURL, secondURL] { + let handle = try FileHandle(forWritingTo: url) + try handle.truncate(atOffset: 0) + try handle.close() + } + + // The checkedEval barrier covers every file. Neither returned array may attempt + // deferred I/O after both are truncated. + assertEqual(try XCTUnwrap(loaded["first"]), firstExpected) + assertEqual(try XCTUnwrap(loaded["second"]), secondExpected) + } + + func testLoadArraysEmpty() throws { + let loaded = try loadArrays(urls: []) + XCTAssertTrue(loaded.isEmpty) + } + + func testLoadArraysMetadataOnlyFile() throws { + let url = try writeRawSafetensor( + header: #"{"__metadata__":{"format":"test"}}"#, + name: "metadata-only.safetensors") + + let (arrays, metadata) = try loadArraysAndMetadata(urls: [url]) + XCTAssertTrue(arrays.isEmpty) + XCTAssertEqual(metadata, ["format": "test"]) + } + + func testLoadArraysFallbackMaterializesEveryTensor() throws { + // MLX accepts gaps, but our ordering parser rejects them and falls back. + let url = try writeRawSafetensor( + header: + #"{"a":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"b":{"dtype":"U8","shape":[1],"data_offsets":[2,3]}}"#, + payload: Data([7, 0, 9]), + name: "fallback.safetensors") + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + + let loaded = try loadArrays(urls: [url]) + XCTAssertTrue(loaded.values.allSatisfy { isEvaluated($0) }) + let handle = try FileHandle(forWritingTo: url) + try handle.truncate(atOffset: 0) + try handle.close() + + XCTAssertEqual(try XCTUnwrap(loaded["a"]).item(UInt8.self), 7) + XCTAssertEqual(try XCTUnwrap(loaded["b"]).item(UInt8.self), 9) + } + + func testDuplicateNameLaterFileWins() throws { + let first = try write(arrays: ["x": MLX.zeros([4])], name: "first.safetensors") + let second = try write(arrays: ["x": MLX.ones([4])], name: "second.safetensors") + + let loaded = try loadArrays(urls: [first, second]) + assertEqual(try XCTUnwrap(loaded["x"]), MLX.ones([4])) + } + + func testMetadataMergedInFileOrder() throws { + let first = try write( + arrays: ["a": MLX.zeros([2])], + metadata: ["shared": "first", "only-first": "1"], + name: "first.safetensors") + let second = try write( + arrays: ["b": MLX.ones([2])], + metadata: ["shared": "second", "only-second": "2"], + name: "second.safetensors") + + let (arrays, metadata) = try loadArraysAndMetadata(urls: [first, second]) + XCTAssertEqual(arrays.keys.sorted(), ["a", "b"]) + XCTAssertEqual( + metadata, ["shared": "second", "only-first": "1", "only-second": "2"]) + } + + func testMissingFileThrows() throws { + let good = try write(arrays: ["a": MLX.zeros([2])], name: "good.safetensors") + let missing = temporaryPath.appending( + path: "missing.safetensors", directoryHint: .notDirectory) + + XCTAssertThrowsError(try loadArrays(urls: [good, missing])) + } + + func testUnknownExtensionThrows() throws { + let bogus = temporaryPath.appending(path: "weights.bin", directoryHint: .notDirectory) + try Data([0, 1, 2, 3]).write(to: bogus) + + XCTAssertThrowsError(try loadArrays(urls: [bogus])) + } + + // MARK: - internals + + func testSafetensorNamesInFileOrder() throws { + let arrays = [ + "big": MLX.zeros([64, 64]), + "small": MLX.zeros([4]), + "medium": MLX.zeros([16, 16]), + ] + let url = try write(arrays: arrays, name: "spans.safetensors") + + let names = try safetensorNamesInFileOrder(url: url) + XCTAssertEqual(names.sorted(), ["big", "medium", "small"]) + } + + func testSafetensorNamesFollowOffsetsRatherThanNameOrHeaderOrder() throws { + let url = try writeRawSafetensor( + header: + #"{"middle":{"dtype":"U8","shape":[2],"data_offsets":[1,3]},"alpha":{"dtype":"U8","shape":[1],"data_offsets":[3,4]},"zeta":{"dtype":"U8","shape":[1],"data_offsets":[0,1]}}"#, + payload: Data([1, 2, 3, 4]), + name: "offset-order.safetensors") + + XCTAssertEqual(try safetensorNamesInFileOrder(url: url), ["zeta", "middle", "alpha"]) + } + + func testSafetensorNamesRejectsMalformedFile() throws { + let url = temporaryPath.appending(path: "bad.safetensors", directoryHint: .notDirectory) + try Data([1, 2, 3]).write(to: url) + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesRejectsOverflowingOffsets() throws { + let url = try writeRawSafetensor( + header: + #"{"bad":{"dtype":"U8","shape":[1],"data_offsets":[-9223372036854775808,9223372036854775807]}}"#, + name: "overflow.safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesRejectsNonintegerOffsets() throws { + for (name, offset) in [("fractional", "1.5"), ("boolean", "true")] { + let url = try writeRawSafetensor( + header: + #"{"bad":{"dtype":"U8","shape":[1],"data_offsets":[0,"# + + offset + #"]}}"#, + name: "\(name).safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + } + + func testSafetensorNamesRejectsNoncontiguousOffsets() throws { + let url = try writeRawSafetensor( + header: + #"{"a":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"b":{"dtype":"U8","shape":[1],"data_offsets":[2,3]}}"#, + payload: Data([0, 0, 0]), + name: "noncontiguous.safetensors") + + XCTAssertThrowsError(try safetensorNamesInFileOrder(url: url)) + } + + func testSafetensorNamesOrdersEmptyTensorBeforeSharedOffset() throws { + let url = try writeRawSafetensor( + header: + #"{"value":{"dtype":"U8","shape":[1],"data_offsets":[0,1]},"empty":{"dtype":"U8","shape":[0],"data_offsets":[0,0]}}"#, + payload: Data([42]), + name: "empty.safetensors") + + XCTAssertEqual(try safetensorNamesInFileOrder(url: url), ["empty", "value"]) + } + + func testEvaluationOrderReversesFilesAndOffsetsWithoutDroppingDuplicateNames() { + let first = ["shared": MLXArray(Int32(1)), "last": MLXArray(Int32(2))] + let second = ["shared": MLXArray(Int32(3)), "last": MLXArray(Int32(4))] + let roots = safetensorArraysInEvaluationOrder( + arrays: [first, second], names: [["shared", "last"], ["shared", "last"]]) + + // eval_impl pops its tape from the back: all four independent Load roots must + // be retained in reverse order, including the name overwritten by the merge. + XCTAssertEqual(roots.map { $0.item(Int32.self) }, [4, 3, 2, 1]) + } + + func testEvaluationOrderIncludesEveryArrayWhenPlanningFailsOrNamesDisagree() { + let arrays = ["a": MLXArray(Int32(1)), "b": MLXArray(Int32(2))] + for names: [String]? in [nil, [], ["a"], ["a", "missing"]] { + let roots = safetensorArraysInEvaluationOrder(arrays: [arrays], names: [names]) + XCTAssertEqual(roots.map { $0.item(Int32.self) }.sorted(), [1, 2]) + } + } +} diff --git a/Tests/MLXTests/MaterializedHazardTests.swift b/Tests/MLXTests/MaterializedHazardTests.swift index 5bc94021..644bb9c6 100644 --- a/Tests/MLXTests/MaterializedHazardTests.swift +++ b/Tests/MLXTests/MaterializedHazardTests.swift @@ -1,5 +1,6 @@ // Copyright © 2026 Apple Inc. +import Cmlx import Foundation @_spi(MaterializedModule) import MLXNN import Testing @@ -336,9 +337,121 @@ private func compileInput() -> MLXArray { MLXArray((0 ..< 8).map { Float($0) }, [2, 4]) } +// Access to the mutable state is handed between threads by the test's semaphores. +private final class ReboundCompileState: @unchecked Sendable { + let holder: ArrayHolder + let compiled: (MLXArray) -> MLXArray + var workerBefore: Float = 0 + var workerAfter: Float = 0 + + init() { + let holder = ArrayHolder(MLXArray(Float(1)).materialized()) + self.holder = holder + self.compiled = compile(inputs: [holder]) { (x: MLXArray) in x * holder.value } + } + + func value() -> Float { + compiled(MLXArray(Float(2))).item(Float.self) + } +} + +private final class ReleasedCompileState: @unchecked Sendable { + var compiled: CompiledFunction? + let functionID: UInt + var before: Float = 0 + var after: Float = 0 + + init() { + let scale = MLXArray(Float(2)).materialized() + let compiled = CompiledFunction(shapeless: false) { [$0[0] * scale] } + self.compiled = compiled + self.functionID = UInt(bitPattern: Unmanaged.passUnretained(compiled).toOpaque()) + } + + // Reuse the old address-derived ID without depending on allocator reuse. + func replacementValue() -> Float { + let result = withEvalLock { + let body = new_mlx_closure { [$0[0] * 3] } + defer { mlx_closure_free(body) } + var replacement = mlx_closure_new() + defer { mlx_closure_free(replacement) } + mlx_detail_compile(&replacement, body, functionID, false, [], 0) + let input = new_mlx_vector_array([MLXArray(Float(2))]) + defer { mlx_vector_array_free(input) } + var output = mlx_vector_array_new() + defer { mlx_vector_array_free(output) } + mlx_closure_apply(&output, replacement, input) + return mlx_vector_array_values(output)[0] + } + return result.item(Float.self) + } +} + @Suite("compile with rebound materialized state") struct CompileReboundMaterializedStateTests { + @Test func deinitErasesCacheOnTheThreadThatCompiled() { + let state = ReleasedCompileState() + let primed = DispatchSemaphore(value: 0) + let released = DispatchSemaphore(value: 0) + let finished = DispatchSemaphore(value: 0) + let worker = Thread { + state.before = state.compiled!.call([MLXArray(Float(2))])[0].item(Float.self) + primed.signal() + if released.wait(timeout: .now() + 30) == .success { + state.after = state.replacementValue() + } + finished.signal() + } + worker.start() + guard primed.wait(timeout: .now() + 30) == .success else { + released.signal() + Issue.record("worker did not populate its compilation cache") + return + } + state.compiled = nil + released.signal() + guard finished.wait(timeout: .now() + 30) == .success else { + Issue.record("worker did not finish the replacement call") + return + } + #expect(state.before == 4) + #expect(state.after == 6) + } + + @Test func rebindingInvalidatesEveryCallingThreadCache() { + let state = ReboundCompileState() + let primed = DispatchSemaphore(value: 0) + let rebound = DispatchSemaphore(value: 0) + let finished = DispatchSemaphore(value: 0) + + #expect(state.value() == 2) + let worker = Thread { + state.workerBefore = state.value() + primed.signal() + if rebound.wait(timeout: .now() + 30) == .success { + state.workerAfter = state.value() + } + finished.signal() + } + worker.start() + guard primed.wait(timeout: .now() + 30) == .success else { + rebound.signal() + Issue.record("worker did not populate its compilation cache") + return + } + + state.holder.value = MLXArray(Float(3)).materialized() + #expect(state.value() == 6) + rebound.signal() + guard finished.wait(timeout: .now() + 30) == .success else { + Issue.record("worker did not finish its cached call") + return + } + #expect(state.workerBefore == 2) + #expect(state.workerAfter == 6) + } + /// Control: plain modules rebound between calls are threaded through as /// inputs and the compiled function follows the rebinding. @Test func plainModuleRebound() { @@ -546,8 +659,40 @@ struct SendableTransformsMaterializedTests { let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in x * scale } - let input = compileInput() - #expect(compiled(input).allClose(input * 3).item(Bool.self)) + for offset in 0 ..< 3 { + let input = compileInput() + Float(offset) + #expect(compiled(input).allClose(input * 3).item(Bool.self)) + #expect(isEvaluated(scale)) + #expect(!hasGraph(scale)) + #expect(scale.item(Float.self) == 3) + } + } + + @Test func sendableCompileCapturingMaterializedArrayConcurrently() async { + let scale = MLXArray(Float(3)).materialized() + let compiled: @Sendable (MLXArray) -> MLXArray = compile { x in + x * scale + } + + await withTaskGroup(of: Bool.self) { group in + for worker in 0 ..< 8 { + group.addTask { + for iteration in 0 ..< 25 { + let input = compileInput() + Float(worker * 25 + iteration) + if !compiled(input).allClose(input * 3).item(Bool.self) { + return false + } + } + return true + } + } + for await ok in group { + #expect(ok) + } + } + #expect(isEvaluated(scale)) + #expect(!hasGraph(scale)) + #expect(scale.item(Float.self) == 3) } /// `CustomFunction` closures may capture a MaterializedArray, including From 6990d618679d15bba674a1fc9fd225feb073cc35 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Wed, 21 Oct 2026 20:14:23 +0400 Subject: [PATCH 02/11] Add byte-progress reporting for safetensors file loading (#427) * Add byte-progress reporting for safetensors file loading Introduces LoadProgress and new loadArrays(url:stream:progressHandler:) and loadArraysAndMetadata(url:stream:progressHandler:) overloads that report bytes read as lazy arrays are evaluated. Uses a custom mlx_io_reader vtable backed by pread() so progress callbacks can be invoked from MLX worker threads. Includes a unit test verifying monotonic progress from 0 to 1. * Add a scoped load-progress handler and fix memory reader SEEK_END `loadWeights()` style helpers -- including the one in mlx-swift-lm -- call the plain `loadArrays(url:)` / `loadArraysAndMetadata(url:)`, so a per-call `progressHandler:` argument can never reach them without changing every caller along the way. Add `withLoadProgressHandler(_:_:)` (sync and async), a task local scoped handler in the style of `withErrorHandler(_:_:)`. The plain file loading functions report byte progress to it when one is installed, so an application can drive a precise model loading progress bar around code it does not own: let container = try await withLoadProgressHandler({ tracker.update($0) }) { try await factory.loadContainer(from: directory, using: tokenizerLoader) } `LoadProgress` gains the `url` of the file being read so progress can be aggregated across the shards of a sharded model. Also fix the SEEK_END case of the in-memory reader, which moved the offset relative to the current position instead of the end of the data. mlx 0.32.1 seeks to the end of the stream to validate the tensor data offsets against the size of the file, so `loadArrays(data:)` would fail there ("The JSON header is N bytes long but the file is only 8 bytes"). Finally, restructure the truncated file test: a truncated file may now be reported either eagerly, while the header is parsed, or lazily, when the arrays are evaluated, and neither happens with the currently vendored mlx/mlx-c. * Update for the mlx-c custom IO reader/writer API ml-explore/mlx-c#130 changed `mlx_io_vtable` so that the callbacks report whether they succeeded: int (*seek)(void*, int64_t off, int whence); size_t (*read)(void*, char* data, size_t n); size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); size_t (*write)(void*, const char* data, size_t n); and `CReader`/`CWriter` now turn a negative seek or a short read/write into a thrown `std::runtime_error`. Bump the mlx-c submodule to that commit (plus the checked-in copies of `io_types.h` and the CMake `GIT_TAG`) and adapt the Swift side: - `FileIOState.seek` returns `0`/`-1` instead of silently ignoring a bad `whence` or a negative resulting offset. - `FileIOState.read` returns the number of bytes actually read, so a file that is truncated after its header was parsed now fails instead of leaving the destination buffer uninitialized. The private implementation is renamed `readBytes` so it cannot be confused with the two public overloads. - The in-memory reader reports an out-of-bounds read as `0` bytes rather than doing nothing, and its `seek` reports failure for an unknown `whence` or a negative offset. - The file reader's `write` reports `0` bytes, so using it as a writer errors out instead of silently discarding the data. This is the error-propagation path the load progress work was waiting on (ml-explore/mlx#3742 is already in the vendored mlx v0.32.2), so the truncated file test asserts a failure rather than skipping, and a new test truncates the file *after* the header is parsed to cover the lazy read path: [mlx_io_reader] unable to read 65536 bytes (read 65504 instead) in file ... Also fix the `withLoadProgressHandler` documentation links -- the `-(_,()throws->R)` disambiguation does not resolve and `verify-docs.sh` builds with `--warnings-as-errors` -- document why the reported progress is approximate (loading is lazy, so it may stop short of the file size and it counts bytes read rather than bytes covered), and add `LoadProgress` to the MLX topics. * Harden progress reads and cover lazy handler lifetime --- Source/MLX/Documentation.docc/MLX.md | 4 + .../MLX/Documentation.docc/free-functions.md | 4 + Source/MLX/IO.swift | 408 ++++++++++++++++++ Tests/MLXTests/SaveTests.swift | 344 +++++++++++++++ 4 files changed, 760 insertions(+) diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md index 46b5ae12..c8e08c22 100644 --- a/Source/MLX/Documentation.docc/MLX.md +++ b/Source/MLX/Documentation.docc/MLX.md @@ -143,3 +143,7 @@ boundaries, take an evaluated, immutable snapshot. ### Distributed - ``MLXDistributed`` + +### Load Progress + +- ``LoadProgress`` diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index e0afc1fe..5fb01d0e 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -120,9 +120,13 @@ operations as methods for convenience. - ``loadArray(url:stream:)`` - ``loadArrays(url:stream:)`` +- ``loadArrays(url:stream:progressHandler:)`` - ``loadArraysAndMetadata(url:stream:)`` +- ``loadArraysAndMetadata(url:stream:progressHandler:)`` - ``save(array:url:stream:)`` - ``save(arrays:metadata:url:stream:)`` +- ``withLoadProgressHandler(_:_:)-3ghip`` +- ``withLoadProgressHandler(_:_:)-8cm31`` ### Logical diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index 98347ade..62bb3b07 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -3,6 +3,109 @@ import Cmlx import Foundation +/// Byte-level progress for loading arrays from disk. +/// +/// `completedUnitCount` and `totalUnitCount` are bytes and describe a single file. +/// Progress callbacks for each load are delivered in monotonically increasing order. +/// Loading the same file again starts a new count. Callbacks for different loads may +/// interleave -- use ``url`` to aggregate progress when loading a model made of several `safetensors` shards. +/// +/// Loading `safetensors` is lazy, which makes the reported progress approximate: +/// +/// - It may stop short of `totalUnitCount`. Only arrays that are actually evaluated are +/// read, so weights that are dropped before evaluation -- the ones a model's +/// `sanitize(weights:metadata:)` discards, for example -- are never read at all. Treat +/// completion of array evaluation as completion rather than waiting for +/// ``fractionCompleted`` to reach `1`. +/// - It counts bytes read rather than bytes of the file covered, so a region that happens +/// to be read more than once is counted more than once. ``completedUnitCount`` is +/// clamped to `totalUnitCount`, so ``fractionCompleted`` never exceeds `1`. +public struct LoadProgress: Sendable, Equatable { + /// The file being read. + public let url: URL + + /// Bytes read so far, clamped to ``totalUnitCount``. + public let completedUnitCount: Int64 + + /// Size of the file in bytes. + public let totalUnitCount: Int64 + + public var fractionCompleted: Double { + guard totalUnitCount > 0 else { return 0 } + return min(1, max(0, Double(completedUnitCount) / Double(totalUnitCount))) + } + + public init(url: URL, completedUnitCount: Int64, totalUnitCount: Int64) { + self.url = url + self.completedUnitCount = completedUnitCount + self.totalUnitCount = totalUnitCount + } +} + +/// Holder for the scoped ``withLoadProgressHandler(_:_:)-3ghip`` handler. +enum LoadProgressHandler { + + /// The stack of installed handlers -- the innermost scope wins. + @TaskLocal + static var handlers: [@Sendable (LoadProgress) -> Void] = [] + + static var current: (@Sendable (LoadProgress) -> Void)? { + handlers.last + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads. +/// +/// Any ``loadArrays(url:stream:)`` or ``loadArraysAndMetadata(url:stream:)`` performed +/// inside `body` reports byte progress to `handler`, without the call site having to pass +/// a progress handler explicitly. This makes it possible to drive a precise loading +/// progress bar for code -- such as a model loading library -- that you do not control: +/// +/// ```swift +/// let tracker = LoadProgressTracker(totalBytes: totalBytesOfSafetensors(in: directory)) +/// let model = try withLoadProgressHandler({ tracker.update($0) }) { +/// try loadModel(from: directory) +/// } +/// ``` +/// +/// Loading is lazy: the reader captures the handler when the file is opened and keeps +/// reporting to it when the arrays are evaluated, even after `body` returns. Include +/// evaluation in `body` to observe all progress before it returns. Arrays that are never +/// evaluated are never read, so progress may legitimately stop short of the file size. +/// +/// - Note: `handler` is called on the loading thread and MLX worker threads, potentially +/// concurrently for different loads, and is on the critical path of the read. It should be cheap and +/// must not call back into MLX loading or evaluation. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +/// +/// ### See Also +/// - ``loadArrays(url:stream:progressHandler:)`` +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () throws -> R +) rethrows -> R { + try LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try body() + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads (async). +/// +/// See ``withLoadProgressHandler(_:_:)-3ghip`` for details. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () async throws -> R +) async rethrows -> R { + try await LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try await body() + } +} + public enum LoadSaveError: Error { case unableToOpen(URL, String) case unknownExtension(String) @@ -117,6 +220,9 @@ public func loadArray(url: URL, stream: StreamOrDevice = .cpu) throws -> MLXArra /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArray(url:stream:)`` /// - ``loadArraysAndMetadata(url:stream:)`` @@ -128,6 +234,10 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArrays(url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -143,12 +253,35 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin } } +/// Load dictionary of ``MLXArray`` from a `safetensors` file, reporting byte progress as +/// lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. Called on the loading thread and MLX worker +/// threads, potentially concurrently across loads. Keep it cheap and do not call +/// MLX loading or evaluation from the callback. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. +public func loadArrays( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> [String: MLXArray] { + let (arrays, _) = try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + return arrays +} + /// Load dictionary of ``MLXArray`` and metadata `[String:String]` from a `safetensors` file. /// /// - Parameters: /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArrays(url:stream:)`` /// - ``loadArray(url:stream:)`` @@ -160,6 +293,11 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -175,6 +313,47 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw } } +/// Load dictionary of ``MLXArray`` and metadata from a `safetensors` file, reporting byte +/// progress as lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. Called on the loading thread and MLX worker +/// threads, potentially concurrently across loads. Keep it cheap and do not call +/// MLX loading or evaluation from the callback. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. +/// +/// ### See Also +/// - ``loadArraysAndMetadata(url:stream:)`` +/// - ``loadArrays(url:stream:progressHandler:)`` +public func loadArraysAndMetadata( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> ([String: MLXArray], [String: String]) { + precondition(url.isFileURL) + + switch url.pathExtension { + case "safetensors": + var r0 = mlx_map_string_to_array_new() + var r1 = mlx_map_string_to_string_new() + defer { mlx_map_string_to_array_free(r0) } + defer { mlx_map_string_to_string_free(r1) } + + let reader = try new_mlx_io_reader_fileIO(url, progressHandler: progressHandler) + defer { mlx_io_reader_free(reader) } + + _ = try withError { + mlx_load_safetensors_reader(&r0, &r1, reader, stream.ctx) + } + + return (mlx_map_array_values(r0), mlx_map_string_values(r1)) + default: + throw LoadSaveError.unknownExtension(url.pathExtension) + } +} + // MARK: - Memory I/O private class IOState { @@ -187,6 +366,190 @@ private class IOState { } } +private final class FileIOState { + private static let maximumReadChunkSize = 4 * 1024 * 1024 + + private let descriptor: CInt + private let lock = NSLock() + private let progressLock = NSLock() + private var offset: Int64 = 0 + private var completedUnitCount: Int64 = 0 + private var readError: String? + private let progressHandler: @Sendable (LoadProgress) -> Void + private let labelPointer: UnsafeMutablePointer + private let url: URL + + let totalUnitCount: Int64 + + init(url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void) throws { + let path = url.path(percentEncoded: false) + let descriptor = path.withCString { open($0, O_RDONLY) } + guard descriptor >= 0 else { + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + var statBuffer = stat() + guard fstat(descriptor, &statBuffer) == 0 else { + let message = String(cString: strerror(errno)) + close(descriptor) + throw LoadSaveError.unableToOpen(url, message) + } + + guard let labelPointer = strdup("file \(path)") else { + close(descriptor) + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + self.descriptor = descriptor + self.totalUnitCount = max(0, Int64(statBuffer.st_size)) + self.progressHandler = progressHandler + self.labelPointer = labelPointer + self.url = url + + progressHandler( + .init(url: url, completedUnitCount: 0, totalUnitCount: totalUnitCount)) + } + + deinit { + close(descriptor) + free(labelPointer) + } + + var isOpen: Bool { + descriptor >= 0 + } + + var good: Bool { + lock.withLock { + readError == nil + } + } + + var label: UnsafePointer { + UnsafePointer(labelPointer) + } + + func tell() -> Int { + lock.withLock { + Int(offset) + } + } + + /// Move the read position. + /// + /// Returns `0` on success and `-1` on failure; `CReader` turns a negative result into + /// a thrown error (ml-explore/mlx-c#130). + @discardableResult + func seek(offset newOffset: Int64, whence: Int32) -> Int32 { + lock.withLock { () -> Int32 in + let updated: Int64 + switch whence { + case SEEK_SET: + updated = newOffset + case SEEK_CUR: + let (position, overflow) = offset.addingReportingOverflow(newOffset) + guard !overflow else { return -1 } + updated = position + case SEEK_END: + // offset is relative to the end of the file, not the current position. + let (position, overflow) = totalUnitCount.addingReportingOverflow(newOffset) + guard !overflow else { return -1 } + updated = position + default: + return -1 + } + guard updated >= 0 else { return -1 } + offset = updated + return 0 + } + } + + /// Read `count` bytes at the current position and advance it. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int) -> Int { + guard let data else { return 0 } + + let readOffset = lock.withLock { + offset + } + let bytesRead = readBytes(to: data, count: count, offset: readOffset) + + lock.withLock { + offset += Int64(bytesRead) + } + + return bytesRead + } + + /// Read `count` bytes at an absolute `offset`, leaving the current position alone. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + /// This is the path used to materialize lazily loaded arrays, so it is called from + /// MLX worker threads. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int, offset readOffset: Int64) -> Int { + guard let data else { return 0 } + + return readBytes(to: data, count: count, offset: readOffset) + } + + private func readBytes( + to data: UnsafeMutablePointer, count: Int, offset readOffset: Int64 + ) -> Int { + guard count >= 0, readOffset >= 0, readOffset <= Int64.max - Int64(count) else { + return 0 + } + + var totalRead = 0 + while totalRead < count { + let chunkSize = min(count - totalRead, Self.maximumReadChunkSize) + let bytesRead = pread( + descriptor, + UnsafeMutableRawPointer(data.advanced(by: totalRead)), + chunkSize, + off_t(readOffset + Int64(totalRead))) + if bytesRead < 0 && errno == EINTR { continue } + guard bytesRead > 0 else { + recordReadError(bytesRead: bytesRead, requestedCount: count - totalRead) + break + } + totalRead += bytesRead + reportProgress(bytesRead: bytesRead) + } + + return totalRead + } + + private func recordReadError(bytesRead: Int, requestedCount: Int) { + let message: String + if bytesRead < 0 { + message = String(cString: strerror(errno)) + } else { + message = "unexpected end of file while reading \(requestedCount) bytes" + } + lock.withLock { + if readError == nil { + readError = message + } + } + } + + private func reportProgress(bytesRead: Int) { + guard bytesRead > 0 else { return } + + progressLock.withLock { + completedUnitCount += min(totalUnitCount - completedUnitCount, Int64(bytesRead)) + let progress = LoadProgress( + url: url, + completedUnitCount: completedUnitCount, + totalUnitCount: totalUnitCount) + progressHandler(progress) + } + } +} + private let label: StaticString = "\0" private func getData(_ writer: mlx_io_writer) -> Data { @@ -291,6 +654,51 @@ private func new_mlx_io_reader_dataIO(_ data: Data) -> mlx_io_reader { return mlx_io_reader_new(ptr, new_mlx_io_vtable_dataIO()) } +private func new_mlx_io_vtable_fileIO() -> mlx_io_vtable { + mlx_io_vtable { ptr in + guard let ptr else { return false } + return Unmanaged.fromOpaque(ptr).takeUnretainedValue().isOpen + } good: { ptr in + guard let ptr else { return false } + let state = Unmanaged.fromOpaque(ptr).takeUnretainedValue() + return state.isOpen && state.good + } tell: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.tell() + + } seek: { ptr, offset, whence in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.seek(offset: Int64(offset), whence: whence) + + } read: { ptr, data, n in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.read(to: data, count: n) + + } read_at_offset: { ptr, data, n, offset in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.read(to: data, count: n, offset: Int64(offset)) + + } write: { _, _, _ in + // read only -- reporting 0 bytes written makes CWriter throw + return 0 + + } label: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.label + + } free: { ptr in + Unmanaged.fromOpaque(ptr!).release() + } +} + +private func new_mlx_io_reader_fileIO( + _ url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> mlx_io_reader { + let ptr = Unmanaged.passRetained(try FileIOState(url: url, progressHandler: progressHandler)) + .toOpaque() + return mlx_io_reader_new(ptr, new_mlx_io_vtable_fileIO()) +} + private func new_mlx_io_writer_dataIO() -> mlx_io_writer { let ptr = Unmanaged.passRetained(IOState()).toOpaque() return mlx_io_writer_new(ptr, new_mlx_io_vtable_dataIO()) diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 0ef8e1d5..273757f7 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -8,6 +8,44 @@ import MLX import XCTest +// All access to the recorded callbacks is protected by the lock. +private final class ProgressRecorder: @unchecked Sendable { + private let lock = NSLock() + private var progress = [LoadProgress]() + + func record(_ progress: LoadProgress) { + lock.withLock { + self.progress.append(progress) + } + } + + var reported: [LoadProgress] { + lock.withLock { progress } + } + + var values: [Double] { + reported.map { $0.fractionCompleted } + } + + /// Fractions reported for a single file, in order. + func values(for url: URL) -> [Double] { + reported.filter { $0.url == url }.map { $0.fractionCompleted } + } + + /// Aggregate fraction across every file seen, by bytes. + var aggregateFraction: Double { + var completed = [URL: Int64]() + var total = [URL: Int64]() + for progress in reported { + completed[progress.url] = progress.completedUnitCount + total[progress.url] = progress.totalUnitCount + } + let totalBytes = total.values.reduce(0, +) + guard totalBytes > 0 else { return 0 } + return Double(completed.values.reduce(0, +)) / Double(totalBytes) + } +} + final class SaveTests: XCTestCase { let temporaryPath = FileManager.default.temporaryDirectory.appending( @@ -46,6 +84,289 @@ final class SaveTests: XCTestCase { assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) } + public func testLoadArraysProgressReportsThroughEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "arrays.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays( + url: safetensorsPath + ) { @Sendable progress in + recorder.record(progress) + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + } + } + + public func testLoadArraysProgressFailsOnTruncatedTensorData() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + var data = try Data(contentsOf: safetensorsPath) + data.removeLast(32) + try data.write(to: safetensorsPath) + + // A truncated file has to be reported either eagerly, while the header is + // parsed (mlx >= 0.32.1 validates the tensor data offsets against the size of + // the file), or lazily, when the arrays are evaluated and the short read is + // turned into an error (ml-explore/mlx-c#130). + do { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } + try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("a truncated safetensors file must not load successfully") + } catch { + // expected + } + } + } + + /// mlx validates the tensor data offsets against the size of the file while it parses + /// the header, so a file that is _already_ truncated fails eagerly. A file that is + /// truncated after the header is parsed can only be caught when the lazy arrays are + /// evaluated and the read comes up short -- `mlx_io_reader` turns that into an error + /// rather than leaving the destination buffer uninitialized (ml-explore/mlx-c#130). + public func testLoadFailsWhenFileIsTruncatedBeforeEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated-late.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + // the header is parsed and validated against the size of the file here, but + // the tensor data is only read once the lazy arrays are evaluated + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { + @Sendable in recorder.record($0) + } + let size = try XCTUnwrap(recorder.reported.first?.totalUnitCount) + + // ... so truncate the file out from under them. `truncate()` shortens the + // inode the reader already has open, unlike a rewrite which may replace it. + XCTAssertEqual( + truncate(safetensorsPath.path(percentEncoded: false), off_t(size - 32)), 0) + + do { + try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("evaluating arrays read from a truncated file must fail") + } catch { + // expected + } + + // the bytes that were read are still accounted for, monotonically, and the + // aggregate stops short of the original size of the file + let fractions = recorder.values + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(fractions.first, 0) + XCTAssertLessThan(try XCTUnwrap(fractions.last), 1) + } + } + + public func testScopedLoadProgressHandler() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "scoped.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + + // note: the plain loadArrays(url:) -- no progress handler passed at the call site + let loadedArrays = try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + return loadedArrays + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(Set(recorder.reported.map(\.url)), [safetensorsPath]) + } + } + + public func testScopedLoadProgressHandlerIsScoped() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "unscoped.safetensors", + directoryHint: .notDirectory + ) + try MLX.save(arrays: ["foo": MLX.ones([128, 128])], url: safetensorsPath) + + let recorder = ProgressRecorder() + withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + } + + // outside the scope nothing is reported + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + + XCTAssertTrue(recorder.reported.isEmpty) + } + } + + public func testScopedLoadProgressAggregatesAcrossFiles() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let shards = try (0 ..< 3).map { index -> URL in + let url = temporaryPath.appending( + path: "shard-\(index).safetensors", + directoryHint: .notDirectory + ) + try MLX.save(arrays: ["w\(index)": MLX.ones([64, 128])], url: url) + return url + } + + let recorder = ProgressRecorder() + try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + // this mimics a model loader: several shards loaded lazily, then evaluated + var weights = [String: MLXArray]() + for url in shards { + let (w, _) = try MLX.loadArraysAndMetadata(url: url) + weights.merge(w) { _, new in new } + } + MLX.eval(Array(weights.values)) + } + + XCTAssertEqual(Set(recorder.reported.map(\.url)), Set(shards)) + for url in shards { + XCTAssertEqual(recorder.values(for: url).last, 1) + } + XCTAssertEqual(recorder.aggregateFraction, 1, accuracy: 1e-9) + } + } + + public func testProgressChunksLargeTensorReads() throws { + let url = temporaryPath.appending(path: "large.safetensors") + let array = MLXArray(Array(repeating: Float(3), count: 3 * 1024 * 1024)) + try MLX.save(arrays: ["large": array], metadata: ["format": "test"], url: url) + + let recorder = ProgressRecorder() + let (loaded, metadata) = try MLX.loadArraysAndMetadata(url: url) { + recorder.record($0) + } + let headerBytes = try XCTUnwrap(recorder.reported.last?.completedUnitCount) + XCTAssertLessThan(headerBytes, Int64(array.nbytes)) + XCTAssertEqual(metadata, ["format": "test"]) + XCTAssertEqual( + try XCTUnwrap(loaded["large"]).asArray(Float.self), array.asArray(Float.self)) + + let reports = recorder.reported + XCTAssertEqual(reports.first?.completedUnitCount, 0) + XCTAssertEqual(reports.last?.fractionCompleted, 1) + XCTAssertEqual( + reports.map(\.completedUnitCount), reports.map(\.completedUnitCount).sorted()) + let tensorReports = reports.filter { $0.completedUnitCount > headerBytes } + XCTAssertGreaterThanOrEqual(tensorReports.count, 3) + for (previous, current) in zip(reports, reports.dropFirst()) { + XCTAssertLessThanOrEqual( + current.completedUnitCount - previous.completedUnitCount, 4 * 1024 * 1024) + } + } + + public func testProgressRetainsHandlerAndSkipsUnevaluatedArrays() throws { + let url = temporaryPath.appending(path: "partial.safetensors") + let values = Array(repeating: Float(7), count: 1024) + try MLX.save(arrays: ["used": MLXArray(values), "unused": MLXArray(values)], url: url) + + let recorder = ProgressRecorder() + let arrays = try withLoadProgressHandler({ recorder.record($0) }) { + try MLX.loadArrays(url: url) + } + let headerBytes = try XCTUnwrap(recorder.reported.last?.completedUnitCount) + XCTAssertEqual(try XCTUnwrap(arrays["used"]).asArray(Float.self), values) + let partial = try XCTUnwrap(recorder.reported.last) + XCTAssertEqual( + partial.completedUnitCount, headerBytes + Int64(values.count * MemoryLayout.size) + ) + XCTAssertLessThan(partial.fractionCompleted, 1) + XCTAssertEqual(try XCTUnwrap(arrays["unused"]).asArray(Float.self), values) + XCTAssertEqual(recorder.reported.last?.fractionCompleted, 1) + } + + public func testNestedProgressHandlersRestoreAfterThrowAndExplicitHandlerWins() throws { + enum TestError: Error { case expected } + let url = temporaryPath.appending(path: "nested.safetensors") + try MLX.save(arrays: ["weight": MLXArray([Float(1), 2])], url: url) + let outer = ProgressRecorder() + let inner = ProgressRecorder() + let explicit = ProgressRecorder() + + try withLoadProgressHandler({ outer.record($0) }) { + XCTAssertThrowsError( + try withLoadProgressHandler({ inner.record($0) }) { + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + throw TestError.expected + }) + XCTAssertTrue(outer.reported.isEmpty) + XCTAssertEqual(inner.reported.last?.fractionCompleted, 1) + + let explicitlyLoaded = try MLX.loadArrays(url: url) { explicit.record($0) } + try checkedEval(Array(explicitlyLoaded.values) as [Any]) + XCTAssertTrue(outer.reported.isEmpty) + XCTAssertEqual(explicit.reported.last?.fractionCompleted, 1) + + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + } + XCTAssertEqual(outer.reported.last?.fractionCompleted, 1) + } + + public func testAsyncProgressScopeInheritsHandler() async throws { + let url = temporaryPath.appending(path: "async.safetensors") + try MLX.save(arrays: ["weight": MLXArray([Float(1), 2])], url: url) + let recorder = ProgressRecorder() + try await withLoadProgressHandler({ recorder.record($0) }) { + try await Task { + let arrays = try MLX.loadArrays(url: url) + try checkedEval(Array(arrays.values) as [Any]) + }.value + } + XCTAssertEqual(recorder.reported.first?.fractionCompleted, 0) + XCTAssertEqual(recorder.reported.last?.fractionCompleted, 1) + } + public func testSaveArray() throws { // single array npy file let path = temporaryPath.appending( @@ -111,4 +432,27 @@ final class SaveTests: XCTestCase { assertEqual(try XCTUnwrap(loaded["big"]), try XCTUnwrap(arrays["big"])) } + /// A truncated in-memory buffer must be reported as an error. The in-memory reader + /// reports a read that runs past the end of the buffer as zero bytes so that + /// `mlx_io_reader` throws instead of leaving the destination uninitialized + /// (ml-explore/mlx-c#130). + public func testLoadFromTruncatedDataFails() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let arrays: [String: MLXArray] = [ + "big": MLX.ones([64, 64]) + ] + + var data = try saveToData(arrays: arrays) + data.removeLast(32) + + do { + let loaded = try loadArrays(data: data) + try checkedEval(Array(loaded.values) as [Any]) + XCTFail("a truncated safetensors buffer must not load successfully") + } catch { + // expected + } + } + } + } From f08a2671c6b805f4aba63d6fa10cd22d0af1d324 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:34:25 -0700 Subject: [PATCH 03/11] turn down warnings from the C++ code (#506) - specifically -Wno-shorten-64-to-32 - reduces build output from 148898 to 5020 lines --- Package.swift | 1 + xcode/xcconfig/Cmlx.xcconfig | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/Package.swift b/Package.swift index 693f05a4..2518c380 100644 --- a/Package.swift +++ b/Package.swift @@ -313,6 +313,7 @@ let cmlx = Target.target( .headerSearchPath("json/single_include/nlohmann"), .headerSearchPath("fmt/include"), .define("MLX_VERSION", to: "\"0.32.0\""), + .disableWarning("shorten-64-to-32"), ], linkerSettings: linkerSettings, plugins: cudaBuildPlugins, diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index 57d2b911..1697339c 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -7,7 +7,7 @@ BUILD_LIBRARY_FOR_DISTRIBUTION = YES HEADER_SEARCH_PATHS = $(SDKROOT)/usr/include/c++/v1 $(SDKROOT)/usr/include $(inherited) $(SRCROOT)/../Source/Cmlx/metal-cpp $(SRCROOT)/../Source/Cmlx/fmt/include $(SRCROOT)/../Source/Cmlx/json/single_include/nlohmann $(SRCROOT)/../Source/Cmlx/mlx/mlx/distributed/jaccl/lib -OTHER_CFLAGS = -isysroot $(SDKROOT) +OTHER_CFLAGS = -isysroot $(SDKROOT) -Wno-shorten-64-to-32 USER_HEADER_SEARCH_PATHS = $(SRCROOT)/../Source/Cmlx/mlx-c $(SRCROOT)/../Source/Cmlx/mlx From ce4684631d840d9cbf8dea636749b8c068e1d6ae Mon Sep 17 00:00:00 2001 From: Hamed Shams Date: Thu, 22 Oct 2026 20:21:32 +0300 Subject: [PATCH 04/11] Fix crossEntropy and nllLoss for a class axis other than the last (#509) Both took the target logits along `axis` but expanded the targets and squeezed the result on axis -1, so any other axis aborted in squeeze with "Cannot squeeze axis -1 with size N". They now expand and squeeze on `axis`, as the Python reference does. Two tests pin the values Python returns for axis 0. Fixes #507 --- Source/MLXNN/Losses.swift | 8 ++++---- Tests/MLXTests/LossTests.swift | 21 +++++++++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/Source/MLXNN/Losses.swift b/Source/MLXNN/Losses.swift index bf8d8542..b6c7c69f 100644 --- a/Source/MLXNN/Losses.swift +++ b/Source/MLXNN/Losses.swift @@ -51,8 +51,8 @@ public func crossEntropy( if targets_as_probs { score = sum(logits * targets, axis: axis) } else { - score = takeAlong(logits, targets.expandedDimensions(axis: -1), axis: axis).squeezed( - axis: -1) + score = takeAlong(logits, targets.expandedDimensions(axis: axis), axis: axis).squeezed( + axis: axis) } let logSumExpLogits = logSumExp(logits, axis: axis) @@ -164,8 +164,8 @@ public func mseLoss(predictions: MLXArray, targets: MLXArray, reduction: LossRed public func nllLoss( inputs: MLXArray, targets: MLXArray, axis: Int = -1, reduction: LossReduction = .none ) -> MLXArray { - let loss = -takeAlong(inputs, targets.expandedDimensions(axis: -1), axis: axis).squeezed( - axis: -1) + let loss = -takeAlong(inputs, targets.expandedDimensions(axis: axis), axis: axis).squeezed( + axis: axis) return reduction.reduce(loss: loss) } diff --git a/Tests/MLXTests/LossTests.swift b/Tests/MLXTests/LossTests.swift index 8aba886a..56616657 100644 --- a/Tests/MLXTests/LossTests.swift +++ b/Tests/MLXTests/LossTests.swift @@ -30,6 +30,27 @@ class LossTests: XCTestCase { XCTAssertTrue(all(isNaN(loss2)).item()) } + func testCrossEntropyAxis() { + // The class axis need not be the last one. Expected values from Python mlx: + // nn.losses.cross_entropy(logits, targets, axis=0) -> [0.169846, 0.169846]. + let logits = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2]) + let targets = MLXArray([0, 1]) + let expected = MLXArray(converting: [0.169846, 0.169846]) + let loss = crossEntropy(logits: logits, targets: targets, axis: 0, reduction: .none) + XCTAssertEqual(loss.shape, [2]) + XCTAssertTrue(allClose(loss, expected, atol: 1e-5).item()) + } + + func testNLLLossAxis() { + // nn.losses.nll_loss(inputs, targets, axis=0) -> [-2, -2] + let inputs = MLXArray(converting: [2.0, -1.0, -1.0, 2.0, 0.0, 0.0], [3, 2]) + let targets = MLXArray([0, 1]) + let expected = MLXArray(converting: [-2.0, -2.0]) + let loss = nllLoss(inputs: inputs, targets: targets, axis: 0, reduction: .none) + XCTAssertEqual(loss.shape, [2]) + XCTAssertTrue(allClose(loss, expected).item()) + } + func testTripletLoss() { // This mainly tests that tripletLoss now works for arbitrary values of p From 43490365d8318644d77e3b9327b198fcb1989fde Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Fri, 23 Oct 2026 10:37:55 -0500 Subject: [PATCH 05/11] [CUDA] Allow selecting nvcc's host compiler, and invalidate on change (#443) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CudaBuild plugin hardwires nvcc's host compiler to the Swift toolchain's own clang++, which cannot be overridden. When that clang is outside the range the installed CUDA release accepts, nvcc refuses to compile anything: crt/host_config.h:151:2: error: -- unsupported clang version! clang version must be less than 21 and greater than 3.2 This is reachable today: Swift 6.3.3 ships clang 21.0.0, while CUDA 13.0 accepts only clang < 21. The plugin passes its resolved path to encuda as --clangpp, which encuda appends as -ccbin= *after* anything in NVCC_PREPEND_FLAGS, so that environment variable cannot win either. Resolve the host compiler from, in priority order: the MLX_CUDA_HOST_COMPILER environment variable, a new `hostCompiler` key in CudaBuild.json, then the toolchain's clang++ as before. Existing configurations are unaffected. It must still be a clang: nvcc preprocesses with the host compiler and SwiftPM compiles the resulting .cpp, so pointing -ccbin at g++ yields a translation unit full of GCC-only `_Float32` declarations that clang then rejects. Also stamp the compilation configuration beside each generated .cpp and treat a changed signature as out of date. encuda's --incremental check compared only input/output modification times, which cannot observe a changed host compiler, std, arch, or include path — so switching host compilers silently reused .cpp files preprocessed by the previous one. That produced a flood of errors from glibc headers rather than anything pointing at the real cause. Co-authored-by: Claude Opus 5 --- Plugins/CudaBuild/plugin.swift | 35 +++++++++++++++++++++++------- Source/Encuda/encuda-compile.swift | 25 +++++++++++++++++++++ 2 files changed, 52 insertions(+), 8 deletions(-) diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift index 15ecd2e8..e078c6af 100644 --- a/Plugins/CudaBuild/plugin.swift +++ b/Plugins/CudaBuild/plugin.swift @@ -16,6 +16,7 @@ struct CudaBuild: BuildToolPlugin { var cppLanguageStandard: String? = nil var verbose: Bool = false var codeGeneration: [CodeGeneration] = [] + var hostCompiler: String? = nil init() {} @@ -28,9 +29,29 @@ struct CudaBuild: BuildToolPlugin { verbose = try c.decodeIfPresent(Bool.self, forKey: .verbose) ?? false codeGeneration = try c.decodeIfPresent([CodeGeneration].self, forKey: .codeGeneration) ?? [] + hostCompiler = try c.decodeIfPresent(String.self, forKey: .hostCompiler) } } + /// Host compiler nvcc should use, in priority order: the + /// `MLX_CUDA_HOST_COMPILER` environment variable, the `hostCompiler` key in + /// CudaBuild.json, then the toolchain's own clang++. + /// + /// It must still be a clang: nvcc preprocesses with it and SwiftPM compiles + /// the generated `.cpp`, so a GCC-preprocessed translation unit will not build. + func resolveHostCompiler(context: PluginContext, settings: Settings) -> String { + if let env = ProcessInfo.processInfo.environment["MLX_CUDA_HOST_COMPILER"], !env.isEmpty { + return env + } + if let configured = settings.hostCompiler, !configured.isEmpty { + return configured + } + guard let clangUrl = try? context.tool(named: "clang++") else { + fatalError("clang++ not found") + } + return clangUrl.url.path + } + func createBuildCommands(context: PluginContext, target: Target) async throws -> [Command] { print("CUDA Build Plugin") @@ -40,12 +61,6 @@ struct CudaBuild: BuildToolPlugin { return [] } - guard let clangUrl = try? context.tool(named: "clang++") else { - fatalError("clang++ not found") - } - - print("Use clang++ at: \(clangUrl.url.path)") - let sourceDir = target.directoryURL print("Source directory: \(sourceDir.path)") @@ -61,6 +76,10 @@ struct CudaBuild: BuildToolPlugin { settings = Settings() } + let hostCompiler = resolveHostCompiler(context: context, settings: settings) + + print("Use clang++ at: \(hostCompiler)") + // Scan source directory for .cu files let sourceDirPath = sourceDir.path.hasSuffix("/") ? sourceDir.path : sourceDir.path + "/" @@ -141,7 +160,7 @@ struct CudaBuild: BuildToolPlugin { "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)", executable: encuda.url, arguments: compileArguments( - clangpp: clangUrl.url.path, + clangpp: hostCompiler, sourceDir: sourceDir.path, verboseFlag: verboseFlag, stdArgs: stdArgs, @@ -171,7 +190,7 @@ struct CudaBuild: BuildToolPlugin { displayName: "Linking CUDA objects", executable: encuda.url, arguments: linkArguments( - clangpp: clangUrl.url.path, + clangpp: hostCompiler, verboseFlag: verboseFlag, stdArgs: stdArgs, incrementalFlag: incrementalFlag, diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift index 34ded3f9..04626004 100644 --- a/Source/Encuda/encuda-compile.swift +++ b/Source/Encuda/encuda-compile.swift @@ -82,12 +82,37 @@ extension Encuda { guard status == 0 else { throw EncudaError.nvccFailed(status) } + writeStamp() #endif } + /// Inputs to the generated `.cpp` that are not themselves input *files*. + /// A modification-time comparison cannot see any of these change, so the + /// signature is stamped beside the output and compared on the next run. + private var configurationSignature: String { + [ + "nvcc=\(nvccPath ?? "")", + "ccbin=\(clangppPath ?? "")", + "std=\(std ?? "")", + "arch=\(ProcessInfo.processInfo.environment["CUDA_ARCH"] ?? "")", + "include=\(includeDirs.joined(separator: ":"))", + ].joined(separator: "\n") + } + + private var stampURL: URL { + URL(fileURLWithPath: output + ".encuda-stamp") + } + + private func writeStamp() { + try? configurationSignature.write(to: stampURL, atomically: true, encoding: .utf8) + } + private func isUpToDate() -> Bool { let fm = FileManager.default let outputURL = URL(fileURLWithPath: output) + guard let stamp = try? String(contentsOf: stampURL, encoding: .utf8), + stamp == configurationSignature + else { return false } guard fm.fileExists(atPath: output), let outputMod = (try? outputURL.resourceValues(forKeys: [.contentModificationDateKey]))? From a8a3d68f69bede2e3548ed7a00b526b8a937ab6f Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Fri, 23 Oct 2026 12:12:31 -0500 Subject: [PATCH 06/11] [CUDA] Make the CUDA include roots configurable (#444) Two include roots are hardcoded or missing in the SwiftPM CUDA build, and neither can be pointed elsewhere. CCCL is hardcoded to the copy inside the CUDA toolkit, for both the `-I` and the `MLX_CCCL_DIR` define that NVRTC uses at runtime. MLX pins its own CCCL version, and when the toolkit ships a different one the mismatch surfaces as overload ambiguity in device code rather than anything naming a version: mlx/backend/cuda/device/binary_ops.cuh(50): error: more than one instance of overloaded function "cuda::std::fmod" matches the argument list Observed with CUDA 13.0 (CCCL 3.0.1) against MLX's pinned CCCL 3.1.3. cudnn-frontend and CUTLASS are not provided at all. The CMake build fetches both via FetchContent, but SwiftPM has no equivalent, so the build fails on headers that are simply absent: mlx/backend/cuda/cudnn_utils.h:9:10: fatal error: 'cudnn_frontend.h' file not found mlx/backend/cuda/quantized/qmm/qmv.cu:9:10: fatal error: 'cute/numeric/numeric_types.hpp' file not found Adds `MLX_CCCL_DIR` to override the CCCL root (defaulting to the toolkit copy, so existing behaviour is unchanged) and `MLX_CUDA_INCLUDE_PATHS`, a colon-separated list of additional roots. Both are applied to the C++ sources via Package.swift and to `.cu` compilation via the CudaBuild plugin, so device and host code compile against the same headers. This keeps the roots external rather than vendoring cudnn-frontend and CUTLASS as submodules; that remains an option if their provisioning should be automatic rather than configured. Co-authored-by: Claude Opus 5 --- Package.swift | 28 +++++++++++++++++++++++----- Plugins/CudaBuild/plugin.swift | 19 +++++++++++++++++++ 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/Package.swift b/Package.swift index 2518c380..07e6f638 100644 --- a/Package.swift +++ b/Package.swift @@ -33,6 +33,23 @@ import PackageDescription let cudaTargets: [Target] = [] #endif +/// Additional include roots for the CUDA build, colon-separated. +/// +/// The CMake build fetches cudnn-frontend and CUTLASS itself, but SwiftPM has no +/// equivalent, so `cudnn_utils.cpp` and the `quantized/qmm` kernels cannot find +/// their headers unless the roots are supplied here. +let extraCudaIncludePaths = + (Context.environment["MLX_CUDA_INCLUDE_PATHS"] ?? "") + .split(separator: ":") + .map(String.init) + .filter { !$0.isEmpty } + +/// CCCL headers used to compile device code and to JIT at runtime. +/// +/// Defaults to the copy shipped with the CUDA toolkit. MLX pins its own CCCL +/// version, so allow pointing at that copy instead when the two differ. +let cudaCcclDir = Context.environment["MLX_CCCL_DIR"] ?? "/usr/local/cuda/include/cccl" + let noMetalCmlxExcludes = [ // Exclude Metal backend files, but keep no_metal.cpp for stubs // "mlx/mlx/backend/metal/no_metal.cpp", @@ -139,11 +156,12 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/quantized/qmm", ] + noMetalCmlxExcludes - cxxSettings = [ - .unsafeFlags(["-I/usr/local/cuda/include"]), - .unsafeFlags(["-I/usr/local/cuda/include/cccl"]), - .define("MLX_CCCL_DIR", to: "\"/usr/local/cuda/include/cccl\""), - ] + cxxSettings = + [ + .unsafeFlags(["-I/usr/local/cuda/include"]), + .unsafeFlags(["-I\(cudaCcclDir)"]), + .define("MLX_CCCL_DIR", to: "\"\(cudaCcclDir)\""), + ] + extraCudaIncludePaths.map { .unsafeFlags(["-I\($0)"]) } linkerSettings = [ .linkedLibrary("gfortran", .when(platforms: [.linux])), diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift index e078c6af..f28c552c 100644 --- a/Plugins/CudaBuild/plugin.swift +++ b/Plugins/CudaBuild/plugin.swift @@ -151,6 +151,22 @@ struct CudaBuild: BuildToolPlugin { } let stdArgs = settings.cppLanguageStandard.map { ["--std", $0] } ?? [] + // Include roots that live outside the package: the CCCL copy to compile + // device code against, plus anything named in MLX_CUDA_INCLUDE_PATHS + // (cudnn-frontend, CUTLASS). Kept in step with Package.swift, which + // applies the same roots to the C++ sources. + let environment = ProcessInfo.processInfo.environment + var externalIncludeRoots: [String] = [] + if let ccclDir = environment["MLX_CCCL_DIR"] { + externalIncludeRoots.append(ccclDir) + } + let extraRoots = environment["MLX_CUDA_INCLUDE_PATHS"] ?? "" + externalIncludeRoots += extraRoots.split(separator: ":").map(String.init) + let externalIncludeArgs: [String] = + externalIncludeRoots + .filter { !$0.isEmpty } + .flatMap { ["-I", $0] } + for inputFile in sourceCuFiles + generatedCuFiles { let outputCpp = URL(string: inputFile.relativePath, relativeTo: outputDir)! .deletingPathExtension().appendingPathExtension("cpp") @@ -166,6 +182,7 @@ struct CudaBuild: BuildToolPlugin { stdArgs: stdArgs, incrementalFlag: incrementalFlag, headerSearchPathArgs: headerSearchPathArgs, + externalIncludeArgs: externalIncludeArgs, input: inputFile.path, output: outputCpp.path ), @@ -214,6 +231,7 @@ struct CudaBuild: BuildToolPlugin { stdArgs: [String], incrementalFlag: [String], headerSearchPathArgs: [String], + externalIncludeArgs: [String], input: String, output: String ) -> [String] { @@ -224,6 +242,7 @@ struct CudaBuild: BuildToolPlugin { arguments += ["--clangpp", clangpp] arguments += ["-I", sourceDir] arguments += headerSearchPathArgs + arguments += externalIncludeArgs arguments += [input] arguments += ["-o", output] return arguments From 5f8f4d11779cf403a690373cea4965a1178c96a2 Mon Sep 17 00:00:00 2001 From: sbkoth Date: Fri, 23 Oct 2026 10:16:16 -0700 Subject: [PATCH 07/11] Add optimizer state save and restore (#487) * Add optimizer state save and restore * Track Adafactor step so compiled checkpoints stay current * Fix optimizer state docs so DocC builds ModuleParameters lives in MLXNN, so a symbol link in MLXOptimizers fails the documentation build. stateSuffixes() also fails once exportState's comment is inherited, so that link is absolute. --- .../Documentation.docc/MLXOptimizers.md | 15 + Source/MLXOptimizers/Optimizers.swift | 292 +++++++++++++++++- Tests/MLXTests/OptimizerTests.swift | 229 ++++++++++++++ 3 files changed, 535 insertions(+), 1 deletion(-) diff --git a/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md b/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md index 02d7345b..23ca9faa 100644 --- a/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md +++ b/Source/MLXOptimizers/Documentation.docc/MLXOptimizers.md @@ -29,6 +29,21 @@ for _ in 0 ..< epochs { } ``` +## Checkpoints + +``Optimizer/state()`` is the buffers training has accumulated (momentum, moments, +and per-parameter step counts), in the same shape as the model parameters. +``Optimizer/update(parameters:)`` writes that tree back. The learning rate and +other hyperparameters stay on the optimizer you construct: + +```swift +eval(optimizer) +let state = optimizer.state() + +let resumed = Adam(learningRate: 1e-3) +resumed.update(parameters: state) +``` + ## Other MLX Packages - [MLX](mlx) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index 42f24751..b66671b5 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -24,6 +24,71 @@ public protocol Optimizer: Updatable, Evaluatable { /// Apply the gradients to the parameters of the model and update the model with the new parameters. func update(model: Module, gradients: ModuleParameters) + + /// Buffers accumulated while training, as a `ModuleParameters` tree. + /// + /// The tree follows the model parameters. Optimizers that keep one buffer per + /// parameter (``SGD``, ``RMSprop``, ``AdaGrad``, ``Lion``, ``Muon``) store that + /// array at the parameter. Optimizers that keep several store a dictionary of + /// named buffers there: + /// + /// - ``Adam`` and ``AdamW``: `m`, `v`, and `step` + /// - ``Adamax``: `m` and `v` + /// - ``AdaDelta``: `v` and `u` + /// - ``Adafactor``: `step`, plus the running averages that parameter uses + /// (`expAvgSq`, or `expAvgSqRow` and `expAvgSqCol`, and `expAvg` when + /// `beta1` is set) + /// + /// ``MultiOptimizer`` nests each sub-optimizer under `states`, indexed + /// `"0"`, `"1"`, ... in optimizer order. + /// `step` is counted per parameter. Hyperparameters such as the learning rate are + /// not part of this tree — build the resumed optimizer with the same values. + /// + /// The result is empty until the first ``update(model:gradients:)`` or + /// ``update(parameters:)``. `flattened()` is the dictionary + /// `save(arrays:metadata:url:)` writes. + /// + /// ```swift + /// eval(optimizer) + /// let checkpoint = optimizer.state() + /// + /// let resumed = Adam(learningRate: 1e-3) + /// resumed.update(parameters: checkpoint) + /// ``` + func state() -> ModuleParameters + + /// Replace buffers with a tree produced by ``state()``. + /// + /// Parameters present in the tree replace the stored buffers. Parameters left + /// out are unchanged, as with `Module.update(parameters:)`. + /// + /// - Parameters: + /// - parameters: buffers from ``state()``, or the same tree read back from + /// `loadArrays(url:)` + func update(parameters: ModuleParameters) +} + +extension Optimizer { + // External `Optimizer` conformers predate this API. The built-in optimizers + // replace both methods; a custom one keeps compiling and returns no buffers + // until it implements them. + public func state() -> ModuleParameters { + ModuleParameters() + } + + public func update(parameters: ModuleParameters) { + } +} + +private func requiredState( + _ arrays: [(String, MLXArray)], keys: [String], owner: Any +) -> [String: MLXArray] { + let byKey = Dictionary(arrays, uniquingKeysWith: { _, latest in latest }) + let missing = keys.filter { byKey[$0] == nil } + precondition( + missing.isEmpty, + "\(type(of: owner)) state is missing \(missing.joined(separator: ", "))") + return byKey } /// An optimizer that delegates to several sub-optimizers, routing each parameter to the first @@ -88,6 +153,38 @@ open class MultiOptimizer: Optimizer { public func innerState() -> [MLXArray] { optimizers.flatMap { $0.innerState() } } + + /// Buffers for every sub-optimizer, under `states.0`, `states.1`, ... + public func state() -> ModuleParameters { + var flat: [(String, MLXArray)] = [] + // `states` matches Python. A bare leading index unflattens as an array, + // which ModuleParameters cannot hold at the root. + for (index, optimizer) in optimizers.enumerated() { + for (key, value) in optimizer.state().flattened() { + flat.append(("states.\(index).\(key)", value)) + } + } + return ModuleParameters.unflattened(flat) + } + + public func update(parameters: ModuleParameters) { + var grouped: [Int: [(String, MLXArray)]] = [:] + for (key, value) in parameters.flattened() { + // "states..", leaving any further dots intact. + let parts = key.split(separator: ".", maxSplits: 2).map(String.init) + guard parts.count == 3, parts[0] == "states", + let index = Int(parts[1]), + optimizers.indices.contains(index) + else { + preconditionFailure( + "MultiOptimizer state key \"\(key)\" is missing a sub-optimizer index") + } + grouped[index, default: []].append((parts[2], value)) + } + for (index, pairs) in grouped { + optimizers[index].update(parameters: ModuleParameters.unflattened(pairs)) + } + } } /// The base class for all optimizers. It allows us to implement an optimizer on a per-parameter basis @@ -124,6 +221,128 @@ open class OptimizerBase: Optimizer { .flatMap { $0.innerState() } } + /// Names written by ``exportState(_:)`` for one parameter. + /// + /// `""` stores the only buffer directly on the parameter. Any other name stores + /// a dictionary of buffers under the parameter (`m`, `v`, `step`, ...). + /// ``stateSuffixes()`` must list every name ``exportState(_:)`` can produce. + open func stateSuffixes() -> [String] { [""] } + + /// Encode one parameter's buffers. + /// + /// Return `[("", array)]` for a single buffer. Return `[(name, array)]` for + /// several, using the names from ``/MLXOptimizers/OptimizerBase/stateSuffixes()``. + open func exportState(_ state: State) -> [(String, MLXArray)] { + fatalError("exportState() not implemented \(type(of: self))") + } + + /// Rebuild one parameter's buffers from the pairs ``exportState(_:)`` produced. + open func importState(_ arrays: [(String, MLXArray)]) -> State { + fatalError("importState() not implemented \(type(of: self))") + } + + open func state() -> ModuleParameters { + ModuleParameters(item: exportTree(stateStorage.asItem())) + } + + open func update(parameters: ModuleParameters) { + let merged = mergeTree(stateStorage.asItem(), parameters.asItem()) + guard case .dictionary(let values) = merged else { + preconditionFailure("optimizer state must be a dictionary") + } + stateStorage = NestedDictionary(values: values) + } + + /// A dictionary of arrays whose keys are all buffer names is one parameter. + /// A dictionary of further dictionaries is the module tree, so a parameter + /// that is itself named `m` or `v` is not mistaken for a buffer. + private func isStateDictionary(_ updates: [String: NestedItem]) -> Bool { + let names = Set(stateSuffixes().filter { !$0.isEmpty }) + guard !names.isEmpty, !updates.isEmpty else { return false } + guard updates.keys.allSatisfy({ names.contains($0) }) else { return false } + for value in updates.values { + guard case .value = value else { return false } + } + return true + } + + private func exportTree(_ item: NestedItem) -> NestedItem { + switch item { + case .none: + return .none + case .value(let state): + return exportValue(state) + case .array(let items): + return .array(items.map { exportTree($0) }) + case .dictionary(let dictionary): + return .dictionary(dictionary.mapValues { exportTree($0) }) + } + } + + private func exportValue(_ state: State) -> NestedItem { + let pairs = exportState(state) + precondition(!pairs.isEmpty, "\(type(of: self)) exportState() returned no buffers") + if pairs.count == 1, pairs[0].0.isEmpty { + return .value(pairs[0].1) + } + + var exported: [String: NestedItem] = [:] + for (name, array) in pairs { + precondition( + !name.isEmpty, "\(type(of: self)) mixed an unnamed buffer with named buffers") + exported[name] = .value(array) + } + return .dictionary(exported) + } + + private func mergeTree( + _ existing: NestedItem, _ update: NestedItem + ) -> NestedItem { + switch update { + case .none: + return existing + case .value(let array): + return .value(importState([("", array)])) + case .array(let updates): + var merged: [NestedItem] + if case .array(let items) = existing { + merged = items + } else { + merged = [] + } + if merged.count < updates.count { + merged.append( + contentsOf: Array(repeating: .none, count: updates.count - merged.count)) + } + for (index, item) in updates.enumerated() { + merged[index] = mergeTree(merged[index], item) + } + return .array(merged) + case .dictionary(let updates): + if isStateDictionary(updates) { + return .value( + importState( + updates.map { name, value in + guard case .value(let array) = value else { + preconditionFailure( + "optimizer state buffer \(name) is not an array") + } + return (name, array) + })) + } + var merged: [String: NestedItem] + if case .dictionary(let items) = existing { + merged = items + } else { + merged = [:] + } + for (key, item) in updates { + merged[key] = mergeTree(merged[key] ?? .none, item) + } + return .dictionary(merged) + } + } + final public func update(model: Module, gradients: ModuleParameters) { model.update(parameters: apply(gradients: gradients, modelParameters: model.parameters())) } @@ -159,6 +378,17 @@ open class OptimizerBaseArrayState: OptimizerBase { override open func newState(parameter: MLXArray) -> MLXArray { MLXArray.zeros(like: parameter) } + + override open func exportState(_ state: MLXArray) -> [(String, MLXArray)] { + [("", state)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> MLXArray { + guard let array = arrays.first(where: { $0.0.isEmpty })?.1 else { + preconditionFailure("\(type(of: self)) state is missing the parameter buffer") + } + return array + } } /// State container for ``OptimizerBase`` holding a tuple of `MLXArray`. @@ -388,6 +618,17 @@ open class AdaDelta: OptimizerBase { TupleState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["v", "u"] } + + override open func exportState(_ state: TupleState) -> [(String, MLXArray)] { + [("v", state.values.0), ("u", state.values.1)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> TupleState { + let byKey = requiredState(arrays, keys: ["v", "u"], owner: self) + return TupleState(byKey["v"]!, byKey["u"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( MLXArray, TupleState ) { @@ -441,6 +682,17 @@ open class Adam: OptimizerBase { AdamState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["m", "v", "step"] } + + override open func exportState(_ state: AdamState) -> [(String, MLXArray)] { + [("m", state.values.0), ("v", state.values.1), ("step", state.step)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> AdamState { + let byKey = requiredState(arrays, keys: ["m", "v", "step"], owner: self) + return AdamState(byKey["m"]!, byKey["v"]!, step: byKey["step"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: AdamState) -> ( MLXArray, AdamState ) { @@ -541,6 +793,17 @@ open class Adamax: OptimizerBase { TupleState(zeros: parameter) } + override open func stateSuffixes() -> [String] { ["m", "v"] } + + override open func exportState(_ state: TupleState) -> [(String, MLXArray)] { + [("m", state.values.0), ("v", state.values.1)] + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> TupleState { + let byKey = requiredState(arrays, keys: ["m", "v"], owner: self) + return TupleState(byKey["m"]!, byKey["v"]!) + } + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( MLXArray, TupleState ) { @@ -643,7 +906,10 @@ open class Adafactor: OptimizerBase { var expAvg: MLXArray? = nil public func innerState() -> [MLXArray] { - [expAvgSqRow, expAvgSqCol, expAvgSq, expAvg].compactMap { $0 } + // `compile` and `eval` only see these arrays. `step` has to be + // here or a compiled run keeps the count from the first trace, + // and `state()` saves that stuck value. + [expAvgSqRow, expAvgSqCol, expAvgSq, expAvg].compactMap { $0 } + [step] } } @@ -696,6 +962,30 @@ open class Adafactor: OptimizerBase { return s } + override open func stateSuffixes() -> [String] { + ["step", "expAvgSqRow", "expAvgSqCol", "expAvgSq", "expAvg"] + } + + override open func exportState(_ state: State) -> [(String, MLXArray)] { + var exported: [(String, MLXArray)] = [("step", state.step)] + if let value = state.expAvgSqRow { exported.append(("expAvgSqRow", value)) } + if let value = state.expAvgSqCol { exported.append(("expAvgSqCol", value)) } + if let value = state.expAvgSq { exported.append(("expAvgSq", value)) } + if let value = state.expAvg { exported.append(("expAvg", value)) } + return exported + } + + override open func importState(_ arrays: [(String, MLXArray)]) -> State { + let byKey = requiredState(arrays, keys: ["step"], owner: self) + var state = State() + state.step = byKey["step"]! + state.expAvgSqRow = byKey["expAvgSqRow"] + state.expAvgSqCol = byKey["expAvgSqCol"] + state.expAvgSq = byKey["expAvgSq"] + state.expAvg = byKey["expAvg"] + return state + } + func rms(_ inputs: MLXArray) -> MLXArray { sqrt(mean(square(inputs))) } diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index 6bc06ff6..08d7d87c 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -436,4 +436,233 @@ class OptimizerTests: XCTestCase { XCTAssertGreaterThan(abs(p2).sum().item(Float.self), 0) } + // MARK: - optimizer state checkpoints + + class StateBlock: Module { + let weight = MLXArray.zeros([3]) + } + + /// Nested parameters, a sibling vector, and an array of parameters. + /// The dotted and indexed keys are what a checkpoint has to round-trip. + class StateNet: Module { + let block = StateBlock() + let bias = MLXArray.zeros([3]) + let heads = [MLXArray.zeros([2]), MLXArray.zeros([2])] + } + + class RankNet: Module { + let vector = MLXArray.zeros([4]) + let matrix = MLXArray.zeros([3, 5]) + } + + func step(_ optimizer: Optimizer, _ model: Module, times: Int) { + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + for _ in 0 ..< times { + optimizer.update(model: model, gradients: gradients) + } + eval(model, optimizer) + } + + /// Value copy. `parameters()` returns the arrays the module will mutate in place. + func detached(_ parameters: ModuleParameters) -> ModuleParameters { + parameters.mapValues { MLXArray($0.asArray(Float.self), $0.shape) } + } + + func parameter(_ parameters: ModuleParameters, _ key: String) -> MLXArray { + parameters.flattened().first { $0.0 == key }!.1 + } + + func assertParametersEqual( + _ lhs: ModuleParameters, _ rhs: ModuleParameters, file: StaticString = #filePath, + line: UInt = #line + ) { + let left = lhs.flattened() + let right = rhs.flattened() + XCTAssertEqual(left.map(\.0), right.map(\.0), file: file, line: line) + for (lhs, rhs) in zip(left, right) { + assertEqual(lhs.1, rhs.1, file: file, line: line) + } + } + + /// Train `stepsBefore + stepsAfter` on one optimizer, and the same trajectory + /// again with the state saved through the flattened dictionary `save` writes. + func assertResumes( + _ make: () -> T, model: () -> Module, stepsBefore: Int = 3, stepsAfter: Int = 1, + file: StaticString = #filePath, line: UInt = #line + ) { + let reference = model() + step(make(), reference, times: stepsBefore + stepsAfter) + let continued = detached(reference.parameters()) + + let trained = model() + let optimizer = make() + step(optimizer, trained, times: stepsBefore) + let encoded = Dictionary(uniqueKeysWithValues: optimizer.state().flattened()) + let decoded = ModuleParameters.unflattened(encoded) + + let resumedModel = model() + resumedModel.update(parameters: detached(trained.parameters())) + let resumed = make() + resumed.update(parameters: decoded) + step(resumed, resumedModel, times: stepsAfter) + + assertParametersEqual( + continued, detached(resumedModel.parameters()), file: file, line: line) + } + + func testStateIsEmptyBeforeTheFirstStep() { + XCTAssertTrue(Adam(learningRate: 0.1).state().isEmpty) + XCTAssertTrue(SGD(learningRate: 0.1, momentum: 0.9).state().isEmpty) + } + + func testStateRoundTripSGD() { + assertResumes({ SGD(learningRate: 0.1, momentum: 0.9) }, model: { StateNet() }) + + // A fresh optimizer from the same weights is not the same trajectory. + // Otherwise the comparison above would pass with the buffers ignored. + let model = StateNet() + step(SGD(learningRate: 0.1, momentum: 0.9), model, times: 3) + let weights = detached(model.parameters()) + let continued = StateNet() + step(SGD(learningRate: 0.1, momentum: 0.9), continued, times: 4) + + let freshModel = StateNet() + freshModel.update(parameters: weights) + step(SGD(learningRate: 0.1, momentum: 0.9), freshModel, times: 1) + assertNotEqual( + parameter(detached(freshModel.parameters()), "bias"), + parameter(detached(continued.parameters()), "bias")) + } + + func testStateRoundTripAdam() { + assertResumes( + { Adam(learningRate: 0.1, biasCorrection: true) }, model: { StateNet() }) + + let model = StateNet() + let optimizer = Adam(learningRate: 0.1, biasCorrection: true) + step(optimizer, model, times: 3) + + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("block.weight.m")) + XCTAssertTrue(keys.contains("block.weight.v")) + XCTAssertTrue(keys.contains("block.weight.step")) + XCTAssertTrue(keys.contains("bias.m")) + XCTAssertTrue(keys.contains("heads.0.m")) + XCTAssertTrue(keys.contains("heads.1.v")) + XCTAssertTrue(keys.contains("heads.1.step")) + + for (key, step) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(step.item(Int32.self), 3, key) + } + + // `state()` on the existential is the real tree, not the empty default. + let existential: Optimizer = optimizer + XCTAssertEqual( + Set(existential.state().flattened().map(\.0)), keys) + + // Omitting `bias` leaves it alone and replaces `block.weight`. + let biasBefore = optimizer.state().flattened().filter { $0.0.hasPrefix("bias.") } + let zeroWeight = optimizer.state().flattened().filter { + $0.0.hasPrefix("block.weight.") + }.map { ($0.0, MLXArray.zeros(like: $0.1)) } + optimizer.update(parameters: ModuleParameters.unflattened(zeroWeight)) + let after = optimizer.state().flattened() + for (key, value) in biasBefore { + assertEqual(after.first { $0.0 == key }!.1, value) + } + for (key, _) in zeroWeight { + let value = after.first { $0.0 == key }!.1 + assertEqual(value, MLXArray.zeros(like: value)) + } + } + + func testStateRoundTripAdaDelta() { + assertResumes({ AdaDelta(learningRate: 0.1) }, model: { StateNet() }) + } + + func testStateRoundTripAdamax() { + assertResumes({ Adamax(learningRate: 0.1) }, model: { StateNet() }) + + let model = TwoParameterModel() + let optimizer = Adamax(learningRate: 0.1) + step(optimizer, model, times: 1) + XCTAssertEqual( + Set(optimizer.state().flattened().map(\.0)), + ["bias.m", "bias.v", "weight.m", "weight.v"]) + } + + func testStateRoundTripAdafactor() { + assertResumes( + { Adafactor(learningRate: 0.1, relativeStep: false) }, model: { RankNet() }) + + let model = RankNet() + let optimizer = Adafactor(learningRate: 0.1, beta1: 0.9, relativeStep: false) + step(optimizer, model, times: 2) + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("vector.expAvgSq")) + XCTAssertTrue(keys.contains("vector.expAvg")) + XCTAssertTrue(keys.contains("vector.step")) + XCTAssertTrue(keys.contains("matrix.expAvgSqRow")) + XCTAssertTrue(keys.contains("matrix.expAvgSqCol")) + XCTAssertFalse(keys.contains("matrix.expAvgSq")) + for (key, step) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(step.item(Int32.self), 2, key) + } + + assertResumes( + { Adafactor(learningRate: 0.1, beta1: 0.9, relativeStep: false) }, + model: { RankNet() }) + } + + /// `step` changes beta2 on every update. compile only writes back arrays + /// from `innerState()`, so a checkpoint after a compiled run must see the + /// same count an eager run would. + func testAdafactorCompiledStepAdvances() { + func parameters(compiled: Bool) -> ModuleParameters { + let model = RankNet() + let optimizer = Adafactor(learningRate: 0.1, warmupInit: true) + let gradients = model.parameters().mapValues { MLXArray.ones(like: $0) } + let token = MLXArray(0) + func once(_ token: MLXArray) -> MLXArray { + optimizer.update(model: model, gradients: gradients) + return token + } + let step = + compiled + ? compile(inputs: [model, optimizer], outputs: [model, optimizer], once) + : once + for _ in 0 ..< 4 { + _ = step(token) + } + eval(model, optimizer) + for (key, value) in optimizer.state().flattened() where key.hasSuffix(".step") { + XCTAssertEqual(value.item(Int32.self), 4, key) + } + return detached(model.parameters()) + } + + assertParametersEqual(parameters(compiled: false), parameters(compiled: true)) + } + + func testStateRoundTripMultiOptimizer() { + func make() -> MultiOptimizer { + MultiOptimizer( + optimizers: [ + SGD(learningRate: 1.0, momentum: 0.5), + Adam(learningRate: 0.1, biasCorrection: true), + ], + filters: [{ key, _ in key == "bias" }]) + } + assertResumes(make, model: { TwoParameterModel() }) + + let model = TwoParameterModel() + let optimizer = make() + step(optimizer, model, times: 2) + let keys = Set(optimizer.state().flattened().map(\.0)) + XCTAssertTrue(keys.contains("states.0.bias")) + XCTAssertTrue(keys.contains("states.1.weight.m")) + XCTAssertTrue(keys.contains("states.1.weight.v")) + XCTAssertTrue(keys.contains("states.1.weight.step")) + } + } From dbce77f5973f9ac9a4039b18e437e9b6c35f8873 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:31:11 -0700 Subject: [PATCH 08/11] mlx v0.32.4 (#508) * mlx v0.32.4 New API - MLXFast.gatedDeltaUpdate - Memory.bufferSize - matrixTransposed - uniqued Changed API - gatherQuantizedMM -- adds globalScale: MLXArray? = nil, biases now has a nil default --- .github/scripts/build-linux-cuda-cmake.sh | 7 - CMakeLists.txt | 23 +- Package.swift | 6 + Source/Cmlx/include-framework/Metal.hpp | 2 +- Source/Cmlx/include-framework/mlx-allocator.h | 56 +- Source/Cmlx/include-framework/mlx-array.h | 80 +- .../mlx-backend-metal-device.h | 11 +- .../mlx-backend-metal-resident.h | 10 +- .../mlx-c-distributed_group.h | 57 +- Source/Cmlx/include-framework/mlx-c-fast.h | 10 + Source/Cmlx/include-framework/mlx-c-memory.h | 1 + Source/Cmlx/include-framework/mlx-c-ops.h | 11 + Source/Cmlx/include-framework/mlx-fast.h | 10 + Source/Cmlx/include-framework/mlx-fft.h | 28 + Source/Cmlx/include-framework/mlx-io-load.h | 86 +- Source/Cmlx/include-framework/mlx-memory.h | 15 +- Source/Cmlx/include-framework/mlx-ops.h | 42 + .../Cmlx/include-framework/mlx-primitives.h | 29 +- Source/Cmlx/include-framework/mlx-stream.h | 4 + .../include-framework/mlx-types-complex.h | 22 +- .../include-framework/mlx-types-half_types.h | 7 - Source/Cmlx/include-framework/mlx-utils.h | 15 +- Source/Cmlx/include-framework/mlx-version.h | 2 +- Source/Cmlx/include/mlx/c/distributed_group.h | 57 +- Source/Cmlx/include/mlx/c/fast.h | 10 + Source/Cmlx/include/mlx/c/memory.h | 1 + Source/Cmlx/include/mlx/c/ops.h | 11 + Source/Cmlx/mlx | 2 +- Source/Cmlx/mlx-c | 2 +- Source/Cmlx/mlx-generated/binary_ops.cpp | 2 +- .../Cmlx/mlx-generated/compiled_preamble.cpp | 75 +- Source/Cmlx/mlx-generated/conv.cpp | 35 +- Source/Cmlx/mlx-generated/cross_entropy.cpp | 152 + .../mlx-generated/cuda/cuda_jit_sources.h | 1116 +++-- Source/Cmlx/mlx-generated/fft.cpp | 24 +- Source/Cmlx/mlx-generated/fp_quantized.cpp | 461 +- .../Cmlx/mlx-generated/fp_quantized_nax.cpp | 377 +- .../mlx-generated/gated_delta_update_nax.cpp | 1968 +++++++++ .../gated_delta_update_nax_vjp.cpp | 3816 +++++++++++++++++ Source/Cmlx/mlx-generated/gemv.cpp | 95 +- Source/Cmlx/mlx-generated/gemv_masked.cpp | 2 +- Source/Cmlx/mlx-generated/masked_scatter.cpp | 2 +- .../Cmlx/mlx-generated/metal/arg_reduce.metal | 24 +- Source/Cmlx/mlx-generated/metal/binary_ops.h | 2 +- Source/Cmlx/mlx-generated/metal/complex.h | 24 +- Source/Cmlx/mlx-generated/metal/conv.metal | 4 +- .../Cmlx/mlx-generated/metal/cross_entropy.h | 133 + Source/Cmlx/mlx-generated/metal/defines.h | 1 + Source/Cmlx/mlx-generated/metal/fp8.h | 10 +- .../Cmlx/mlx-generated/metal/fp_quantized.h | 451 +- .../mlx-generated/metal/fp_quantized_nax.h | 367 +- .../mlx-generated/metal/gated_delta_nax_ops.h | 588 +++ .../mlx-generated/metal/gated_delta_update.h | 366 ++ .../metal/gated_delta_update.metal | 50 + .../metal/gated_delta_update_nax.h | 251 ++ .../metal/gated_delta_update_nax_vjp.h | 1737 ++++++++ .../metal/gated_delta_update_vjp.h | 184 + .../metal/gated_delta_update_vjp.metal | 35 + .../metal/gather_mm_offsets.metal | 23 + Source/Cmlx/mlx-generated/metal/gemv_masked.h | 2 +- .../metal/indexing/masked_scatter.h | 2 +- .../mlx-generated/metal/indexing/scatter.h | 6 +- Source/Cmlx/mlx-generated/metal/quantized.h | 462 +- .../Cmlx/mlx-generated/metal/quantized_nax.h | 437 +- .../metal/reduction/reduce_col.h | 14 +- .../metal/reduction/reduce_init.h | 5 +- .../metal/scaled_dot_product_attention.metal | 60 +- .../scaled_dot_product_attention_vjp.metal | 154 + .../Cmlx/mlx-generated/metal/sdpa_blocked.h | 45 + Source/Cmlx/mlx-generated/metal/sdpa_vector.h | 97 +- .../steel/attn/kernels/steel_attention.h | 148 +- .../steel/attn/kernels/steel_attention.metal | 45 +- .../steel/attn/kernels/steel_attention_nax.h | 127 +- .../metal/steel/conv/kernels/steel_conv.h | 4 +- .../steel/conv/kernels/steel_conv_general.h | 2 +- .../steel/conv/loaders/loader_channel_l.h | 28 +- .../steel/conv/loaders/loader_channel_n.h | 7 +- .../metal/steel/conv/loaders/loader_general.h | 8 +- .../steel/gemm/kernels/steel_gemm_gather.h | 283 +- .../gemm/kernels/steel_gemm_gather_nax.h | 142 +- .../steel/gemm/kernels/steel_gemm_segmented.h | 4 +- .../gemm/kernels/steel_gemm_segmented_nax.h | 4 +- .../steel/gemm/kernels/steel_gemm_thin_nax.h | 223 + Source/Cmlx/mlx-generated/metal/unary_ops.h | 10 +- Source/Cmlx/mlx-generated/metal/utils.h | 70 + Source/Cmlx/mlx-generated/quantized.cpp | 462 +- Source/Cmlx/mlx-generated/quantized_nax.cpp | 437 +- Source/Cmlx/mlx-generated/reduce.cpp | 19 +- Source/Cmlx/mlx-generated/scan.cpp | 2 +- Source/Cmlx/mlx-generated/scatter.cpp | 6 +- Source/Cmlx/mlx-generated/steel_attention.cpp | 146 +- .../mlx-generated/steel_attention_nax.cpp | 127 +- Source/Cmlx/mlx-generated/steel_conv.cpp | 4 +- .../Cmlx/mlx-generated/steel_conv_general.cpp | 10 +- .../Cmlx/mlx-generated/steel_gemm_gather.cpp | 283 +- .../mlx-generated/steel_gemm_gather_nax.cpp | 144 +- .../mlx-generated/steel_gemm_segmented.cpp | 4 +- .../steel_gemm_segmented_nax.cpp | 4 +- .../mlx-generated/steel_gemm_thin_nax.cpp | 242 ++ Source/Cmlx/mlx-generated/unary_ops.cpp | 20 +- Source/Cmlx/mlx-generated/utils.cpp | 95 +- Source/MLX/Distributed.swift | 8 +- .../Organization/arithmetic.md | 2 +- Source/MLX/MLXFast.swift | 36 + Source/MLX/Memory.swift | 17 + Source/MLX/Ops+Array.swift | 16 + Source/MLX/Ops.swift | 185 +- .../Generated/GeneratedFastTests.swift | 276 +- .../GeneratedQuantizationTests.swift | 116 +- .../Generated/GeneratedShapeTests.swift | 444 +- Tests/MLXTests/FP8Tests.swift | 20 +- Tests/MLXTests/MemoryTests.swift | 50 + cmake/mlx.patch | 34 - tools/fix-metal-includes.sh | 4 + tools/integration_tests/README.md | 11 +- tools/integration_tests/cases.py | 255 ++ xcode/MLX.xcodeproj/project.pbxproj | 71 +- 117 files changed, 16103 insertions(+), 2828 deletions(-) create mode 100644 Source/Cmlx/mlx-generated/cross_entropy.cpp create mode 100644 Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp create mode 100644 Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp create mode 100644 Source/Cmlx/mlx-generated/metal/cross_entropy.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update.metal create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h create mode 100644 Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal create mode 100644 Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal create mode 100644 Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal create mode 100644 Source/Cmlx/mlx-generated/metal/sdpa_blocked.h create mode 100644 Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h create mode 100644 Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp delete mode 100644 cmake/mlx.patch diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index 2fc81d58..1262ded0 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.github/scripts/build-linux-cuda-cmake.sh @@ -17,13 +17,6 @@ mkdir -p build pushd build cmake -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=ON -DMLX_C_BUILD_EXAMPLES=OFF .. -G Ninja -# Verify the local mlx patch is actually in place: without it the examples below -# crash on exit, which is easy to misread as an unrelated CUDA failure. -git -C _deps/mlx-src apply --reverse --check "$PWD/../cmake/mlx.patch" || { - echo "error: cmake/mlx.patch is not applied to _deps/mlx-src" - exit 1 -} - ninja ./example1 --device gpu ./tutorial --device gpu diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c558a90..cf07480a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -37,22 +37,10 @@ endif() set(apply_patch ${CMAKE_COMMAND} -DREPO=) set(apply_patch_script -P ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake) -# two local patches for CUDA: -# -# * https://github.com/ml-explore/mlx/pull/4480: leak the global command encoder -# map rather than synchronizing on process shutdown -# * make ~CudaHandle non-throwing: cu::Worker is detached and holds the last -# reference to itself, so ~Worker -> ~CudaEvent -> cudaEventDestroy runs on -# the worker thread at an arbitrary point during exit, racing the CUDA -# runtime's own teardown. A failing destroy there terminated the process. -set(mlx_patch - ${apply_patch} -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/mlx.patch - ${apply_patch_script}) - # note: this must be declared (and named `mlx`) before mlx-c is made available -# so that our pinned/patched version wins over the one mlx-c declares. Because -# we now pull mlx in first, we are also responsible for the option defaults that -# mlx-c would normally set for us. +# so that our pinned version wins over the one mlx-c declares. Because we now +# pull mlx in first, we are also responsible for the option defaults that mlx-c +# would normally set for us. set(MLX_BUILD_TESTS OFF) set(MLX_BUILD_EXAMPLES OFF) set(MLX_BUILD_BENCHMARKS OFF) @@ -61,8 +49,7 @@ set(MLX_BUILD_PYTHON_BINDINGS OFF) FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG "v0.32.2" - PATCH_COMMAND ${mlx_patch}) + GIT_TAG "v0.32.4") FetchContent_MakeAvailable(mlx) # mlx-c @@ -70,7 +57,7 @@ set(MLX_C_BUILD_EXAMPLES OFF) FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" - GIT_TAG "ebc88f10caa1b625e6b581437a8dea6df8a70085") + GIT_TAG "19d3c9c180dddf51be441b55e7f454861899bc93") FetchContent_MakeAvailable(mlx-c) # swift-numerics diff --git a/Package.swift b/Package.swift index 07e6f638..7854c2fa 100644 --- a/Package.swift +++ b/Package.swift @@ -58,6 +58,7 @@ let noMetalCmlxExcludes = [ "mlx/mlx/backend/metal/compiled.cpp", "mlx/mlx/backend/metal/conv.cpp", "mlx/mlx/backend/metal/copy.cpp", + "mxl/mlx/backend/metal/cross_entropy.cpp", "mlx/mlx/backend/metal/custom_kernel.cpp", "mlx/mlx/backend/metal/device.cpp", "mlx/mlx/backend/metal/device_info.cpp", @@ -66,6 +67,7 @@ let noMetalCmlxExcludes = [ "mlx/mlx/backend/metal/event.cpp", "mlx/mlx/backend/metal/fence.cpp", "mlx/mlx/backend/metal/fft.cpp", + "mlx/mlx/backend/metal/gated_delta_update.cpp", "mlx/mlx/backend/metal/hadamard.cpp", "mlx/mlx/backend/metal/indexing.cpp", "mlx/mlx/backend/metal/jit_kernels.cpp", @@ -98,6 +100,7 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/conv.cpp", "mlx/mlx/backend/cuda/cublas_utils.cpp", "mlx/mlx/backend/cuda/cudnn_utils.cpp", + "mlx/mlx/backend/cuda/cusolver_utils.cpp", "mlx/mlx/backend/cuda/custom_kernel.cpp", "mlx/mlx/backend/cuda/delayload.cpp", "mlx/mlx/backend/cuda/device_info.cpp", @@ -114,6 +117,7 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp", "mlx/mlx/backend/cuda/slicing.cpp", "mlx/mlx/backend/cuda/utils.cpp", + "mlx/mlx/backend/cuda/wddm.cpp", "mlx/mlx/backend/cuda/worker.cpp", "mlx/mlx/backend/cuda/binary", "mlx/mlx/backend/cuda/conv", @@ -151,6 +155,7 @@ let noCudaCmlxExcludes = [ "mlx-c/mlx/c/metal.cpp", "mlx/mlx/backend/cuda/delayload.cpp", // For Windows + "mlx/mlx/backend/cuda/wddm.cpp", // built by the CudaBuild plugin (nvcc), not by SwiftPM "mlx/mlx/backend/cuda/quantized/qmm", @@ -170,6 +175,7 @@ let noCudaCmlxExcludes = [ .linkedLibrary("openblas", .when(platforms: [.linux])), .unsafeFlags(["-L/usr/local/cuda/lib64"]), .unsafeFlags(["-L/usr/local/cuda/lib64/stubs"]), + .linkedLibrary("cusolver"), .linkedLibrary("cudnn"), .linkedLibrary("cublas"), .linkedLibrary("cublasLt"), diff --git a/Source/Cmlx/include-framework/Metal.hpp b/Source/Cmlx/include-framework/Metal.hpp index a44ec566..f7dbd5ce 100644 --- a/Source/Cmlx/include-framework/Metal.hpp +++ b/Source/Cmlx/include-framework/Metal.hpp @@ -2,7 +2,7 @@ // // Metal.hpp.in // -// Autogenerated from commit 8f784840de21cacd0c5c237ea860179e47dcec91. +// Autogenerated from commit 26091dacff5333233dda0be4bce5297fb3d18ec9. // // Copyright 2020-2024 Apple Inc. // diff --git a/Source/Cmlx/include-framework/mlx-allocator.h b/Source/Cmlx/include-framework/mlx-allocator.h index b014aee0..9823a63f 100644 --- a/Source/Cmlx/include-framework/mlx-allocator.h +++ b/Source/Cmlx/include-framework/mlx-allocator.h @@ -4,9 +4,20 @@ #pragma once #include +#include +#include #include +namespace mlx::core::allocator { +class Data; +} // namespace mlx::core::allocator + +namespace mlx::core::cu { +class CommandEncoder; +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); +} // namespace mlx::core::cu + namespace mlx::core::allocator { // Simple wrapper around buffer pointers @@ -52,12 +63,47 @@ class MLX_API Allocator { MLX_API Allocator& allocator(); -inline Buffer malloc(size_t size) { - return allocator().malloc(size); -} +using Deleter = std::function; + +// The allocator owns data from malloc +class Data { + public: + Data(Buffer buffer, Deleter d) : buffer_(buffer), deleter_(std::move(d)) {} + Data(const Data& other) = delete; + Data& operator=(const Data& other) = delete; + Data(Data&& other) noexcept + : buffer_(std::exchange(other.buffer_, Buffer(nullptr))), + deleter_(std::exchange(other.deleter_, nullptr)), + owned_(std::exchange(other.owned_, false)) {} + ~Data() { + if (owned_) { + allocator().free(buffer_); + } else if (deleter_) { + deleter_(buffer_); + } + } + + Buffer buffer() const& { + return buffer_; + } + Buffer buffer() const&& = delete; + + bool is_owned() const { + return owned_; + } + + private: + Buffer buffer_; + Deleter deleter_; + bool owned_{false}; + + explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {} + friend Data malloc(size_t size); + friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder); +}; -inline void free(Buffer buffer) { - allocator().free(buffer); +inline Data malloc(size_t size) { + return Data(allocator().malloc(size)); } // Make a Buffer from a raw pointer of the given size without a copy. If a diff --git a/Source/Cmlx/include-framework/mlx-array.h b/Source/Cmlx/include-framework/mlx-array.h index 9c6819c2..5fde4ee6 100644 --- a/Source/Cmlx/include-framework/mlx-array.h +++ b/Source/Cmlx/include-framework/mlx-array.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -19,7 +20,6 @@ namespace mlx::core { // Forward declaration class Primitive; -using Deleter = std::function; using ShapeElem = int32_t; using Shape = SmallVector; using Strides = SmallVector; @@ -69,25 +69,30 @@ class MLX_API array { Dtype dtype, const std::function& deleter); - /* Build an array from a buffer */ - explicit array( - allocator::Buffer data, - Shape shape, - Dtype dtype, - Deleter deleter = allocator::free); + /* Build an array from data */ + explicit array(allocator::Data data, Shape shape, Dtype dtype); /** Assignment to rvalue does not compile. */ array& operator=(const array& other) && = delete; array& operator=(array&& other) && = delete; /** Default copy and move constructors otherwise. */ - array& operator=(array&& other) & = default; array(const array& other) = default; array(array&& other) = default; + /** Assignment goes through release() to break sibling reference cycle. */ array& operator=(const array& other) & { - if (this->id() != other.id()) { - this->array_desc_ = other.array_desc_; + if (array_desc_ != other.array_desc_) { + // Hold the old descriptor while assigning, and then release it. Otherwise + // when |this| and |other| are siblings they would be both destroyed. + release(std::exchange(array_desc_, other.array_desc_)); + } + return *this; + } + + array& operator=(array&& other) & noexcept { + if (array_desc_ != other.array_desc_) { + release(std::exchange(array_desc_, std::move(other.array_desc_))); } return *this; } @@ -229,22 +234,7 @@ class MLX_API array { return reinterpret_cast(array_desc_->primitive.get()); } - struct Data { - allocator::Buffer buffer; - Deleter d; - Data(allocator::Buffer buffer, Deleter d = allocator::free) - : buffer(buffer), d(d) {} - // Not copyable - Data(const Data& d) = delete; - Data& operator=(const Data& d) = delete; - Data(Data&& o) : buffer(o.buffer), d(o.d) { - o.buffer = allocator::Buffer(nullptr); - o.d = [](allocator::Buffer) {}; - } - ~Data() { - d(buffer); - } - }; + using Data = allocator::Data; struct Flags { // True iff there are no gaps in the underlying data. Each item @@ -293,7 +283,8 @@ class MLX_API array { /** True indicates the arrays buffer is safe to reuse */ bool is_donatable() const { - return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1); + return array_desc_.use_count() == 1 && + (array_desc_->data.use_count() == 1) && array_desc_->data->is_owned(); } /** The array's siblings. */ @@ -350,11 +341,8 @@ class MLX_API array { return array_desc_->data_size; } - allocator::Buffer& buffer() { - return array_desc_->data->buffer; - } - const allocator::Buffer& buffer() const { - return array_desc_->data->buffer; + allocator::Buffer buffer() const { + return array_desc_->data->buffer(); } size_t buffer_size() const { @@ -438,24 +426,14 @@ class MLX_API array { // Check if the array is a tracer array bool is_tracer() const; - void set_data(allocator::Buffer buffer, Deleter d = allocator::free); + void set_data(allocator::Data data); void set_data( - allocator::Buffer buffer, + allocator::Data data, size_t data_size, Strides strides, Flags flags, - Deleter d) { - set_data(buffer, data_size, std::move(strides), flags, 0, std::move(d)); - } - - void set_data( - allocator::Buffer buffer, - size_t data_size, - Strides strides, - Flags flags, - int64_t offset = 0, - Deleter d = allocator::free); + int64_t offset = 0); void copy_shared_buffer( const array& other, @@ -467,10 +445,12 @@ class MLX_API array { void copy_shared_buffer(const array& other); void overwrite_descriptor(const array& other) { - array_desc_ = other.array_desc_; + release(std::exchange(array_desc_, other.array_desc_)); } - ~array(); + ~array() { + release(std::move(array_desc_)); + } private: // Initialize the arrays data @@ -533,6 +513,12 @@ class MLX_API array { // the primitive which knows how to compute the array's data from its inputs // and the list of array's inputs for the primitive. std::shared_ptr array_desc_; + + /** + * Lets go of an ArrayDesc, breaking the sibling reference cycle of a + * multi-output primitive. + */ + static void release(std::shared_ptr desc); }; template diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-device.h b/Source/Cmlx/include-framework/mlx-backend-metal-device.h index b7c06c1a..3521bbee 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-device.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-device.h @@ -63,8 +63,13 @@ class MLX_API CommandEncoder { template >> void set_vector_bytes(const Vec& vec, size_t nelems, int idx) { - get_command_encoder()->setBytes( - vec.data(), nelems * sizeof(typename Vec::value_type), idx); + using T = typename Vec::value_type; + if (nelems > 0) { + get_command_encoder()->setBytes(vec.data(), nelems * sizeof(T), idx); + } else { + T val{}; + get_command_encoder()->setBytes(&val, sizeof(T), idx); + } } template >> void set_vector_bytes(const Vec& vec, int idx) { @@ -90,7 +95,7 @@ class MLX_API CommandEncoder { } void barrier(); - void end_encoding(); + void end_encoding(bool wait_in_flight = false); void wait_event(Event event, uint64_t value); void signal_event(Event event, uint64_t value); bool needs_commit() const; diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-resident.h b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h index a5c6b58f..a72a1cfe 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-resident.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h @@ -1,11 +1,14 @@ #ifdef __cplusplus -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #pragma once #include +#include #include #include +#include +#include #include #include @@ -97,7 +100,7 @@ class ResidencySets { // gave us a set; the new set is the last one. add_to_set_locked returns the // set it used. Neither it nor remove_from_set_locked commits, so a bulk // resize costs one commit per set instead of one per allocation. - bool add_set_locked(NS::Error** error = nullptr); + bool add_set_locked(std::string* error = nullptr); uint32_t choose_set_locked(size_t bytes); uint32_t add_to_set_locked(const MTL::Allocation* buf, Placement& at); void remove_from_set_locked(const MTL::Allocation* buf, Placement& at); @@ -119,6 +122,9 @@ class ResidencySets { // taking the lock. std::atomic num_sets_{0}; mutable std::mutex mtx_; + std::condition_variable refresh_cv_; + bool stop_refresh_{false}; + std::thread refresh_thread_; }; } // namespace mlx::core::metal diff --git a/Source/Cmlx/include-framework/mlx-c-distributed_group.h b/Source/Cmlx/include-framework/mlx-c-distributed_group.h index c55ee98c..d79b6b22 100644 --- a/Source/Cmlx/include-framework/mlx-c-distributed_group.h +++ b/Source/Cmlx/include-framework/mlx-c-distributed_group.h @@ -4,6 +4,7 @@ #define MLX_DISTRIBUTED_GROUP_H #include +#include #include @@ -41,15 +42,67 @@ int mlx_distributed_init( bool strict, const char* bk /* may be null */); +/** + * Byte-level all-gather callback used by the JACCL side channel. + * + * \param rank The rank of the process. + * \param size The size of the group. + * \param src Pointer to this rank's data of size n_bytes. + * \param dst Pointer to an output buffer of size size * n_bytes. After the + * call, dst[rank * n_bytes, (rank+1) * n_bytes] contains the data from rank. + * \param n_bytes The number of bytes contributed by each rank. + * \param payload Opaque user data passed to the callback. + * \return Zero on success, non-zero on failure. + */ +typedef int (*mlx_distributed_all_gather_fn)( + int rank, + int size, + const char* src, + char* dst, + size_t n_bytes, + void* payload); + +/** + * Initialize JACCL with a custom metadata all-gather callback. + * + * `bk` must be `jaccl`. The synchronous callback reads `n_bytes` from `src` + * and fills `size * n_bytes` in `dst`, ordered by `rank`. Buffers are borrowed. + * The callback returns zero on success, nonzero on failure. + * + * If a `jaccl` group is already cached, it is returned and `all_gather` is + * never called: the callback is dropped and, with `dtor`, `payload` is + * destroyed before this function returns. Call + * `mlx_distributed_clear_backends()` first to initialize with a new callback. + * + * With `dtor`, ownership of `payload` transfers on entry, including failure. + * Destruction occurs once after all references are released. Cached or + * live groups may delay it. Without `dtor`, `payload` is borrowed and must + * remain valid while the callback is retained. + * + * Returns zero on success, nonzero on error. `res` is unchanged on failure. + */ +int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)); + +/** + * Clear the distributed backend cache. + */ +int mlx_distributed_clear_backends(void); + /** * Get the rank. */ -int mlx_distributed_group_rank(mlx_distributed_group group); +int mlx_distributed_group_get_rank(int* res, mlx_distributed_group group); /** * Get the group size. */ -int mlx_distributed_group_size(mlx_distributed_group group); +int mlx_distributed_group_get_size(int* res, mlx_distributed_group group); /** * Split the group. diff --git a/Source/Cmlx/include-framework/mlx-c-fast.h b/Source/Cmlx/include-framework/mlx-c-fast.h index eac10b39..4763516c 100644 --- a/Source/Cmlx/include-framework/mlx-c-fast.h +++ b/Source/Cmlx/include-framework/mlx-c-fast.h @@ -96,6 +96,16 @@ int mlx_fast_cuda_kernel_apply( const mlx_fast_cuda_kernel_config config, const mlx_stream stream); +int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s); int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, diff --git a/Source/Cmlx/include-framework/mlx-c-memory.h b/Source/Cmlx/include-framework/mlx-c-memory.h index 8ae92ae4..839f1906 100644 --- a/Source/Cmlx/include-framework/mlx-c-memory.h +++ b/Source/Cmlx/include-framework/mlx-c-memory.h @@ -30,6 +30,7 @@ extern "C" { int mlx_clear_cache(void); int mlx_get_active_memory(size_t* res); +int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays); int mlx_get_cache_memory(size_t* res); int mlx_get_memory_limit(size_t* res); int mlx_get_peak_memory(size_t* res); diff --git a/Source/Cmlx/include-framework/mlx-c-ops.h b/Source/Cmlx/include-framework/mlx-c-ops.h index 7531fe08..0e32d381 100644 --- a/Source/Cmlx/include-framework/mlx-c-ops.h +++ b/Source/Cmlx/include-framework/mlx-c-ops.h @@ -548,6 +548,7 @@ int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s); int mlx_gather_qqmm( @@ -703,6 +704,7 @@ int mlx_matmul( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -1360,6 +1362,15 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s); int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/include-framework/mlx-fast.h b/Source/Cmlx/include-framework/mlx-fast.h index 345f9fd6..e8b6ec06 100644 --- a/Source/Cmlx/include-framework/mlx-fast.h +++ b/Source/Cmlx/include-framework/mlx-fast.h @@ -61,6 +61,16 @@ MLX_API array scaled_dot_product_attention( bool force_fused = false, StreamOrDevice s = {}); +MLX_API std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state = std::nullopt, + const std::optional& mask = std::nullopt, + StreamOrDevice s = {}); + using TemplateArg = std::variant; using ScalarArg = std::variant; diff --git a/Source/Cmlx/include-framework/mlx-fft.h b/Source/Cmlx/include-framework/mlx-fft.h index 62811390..5b7cc2fe 100644 --- a/Source/Cmlx/include-framework/mlx-fft.h +++ b/Source/Cmlx/include-framework/mlx-fft.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include @@ -213,6 +214,7 @@ inline array irfft2( StreamOrDevice s = {}) { return irfftn(a, axes, norm, s); } + /** Compute the discrete Fourier Transform sample frequencies. */ MLX_API array fftfreq(int n, double d = 1.0, StreamOrDevice s = {}); @@ -235,5 +237,31 @@ MLX_API array ifftshift(const array& a, StreamOrDevice s = {}); MLX_API array ifftshift(const array& a, const std::vector& axes, StreamOrDevice s = {}); +/** Compute the Short-Time Fourier Transform of the last axis of `x`. */ +MLX_API array stft( + const array& x, + int n_fft = 2048, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + const std::string& pad_mode = "reflect", + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + StreamOrDevice s = {}); + +/** Compute the inverse Short-Time Fourier Transform. */ +MLX_API array istft( + const array& stft_matrix, + const std::optional& n_fft = std::nullopt, + const std::optional& hop_length = std::nullopt, + const std::optional& win_length = std::nullopt, + const std::optional& window = std::nullopt, + bool center = true, + FFTNorm norm = FFTNorm::Backward, + const std::optional& onesided = std::nullopt, + const std::optional& length = std::nullopt, + StreamOrDevice s = {}); + } // namespace mlx::core::fft #endif diff --git a/Source/Cmlx/include-framework/mlx-io-load.h b/Source/Cmlx/include-framework/mlx-io-load.h index fc9a942a..c97010d9 100644 --- a/Source/Cmlx/include-framework/mlx-io-load.h +++ b/Source/Cmlx/include-framework/mlx-io-load.h @@ -4,10 +4,12 @@ #pragma once #include +#include #include +#include #include -#ifdef _MSC_VER +#ifdef _WIN32 #include #else #include @@ -54,6 +56,8 @@ class Writer { std::ios_base::seekdir way = std::ios_base::beg) = 0; virtual void write(const char* data, size_t n) = 0; virtual std::string label() const = 0; + virtual void open() {} + virtual ~Writer() = default; }; @@ -76,20 +80,30 @@ class ParallelFileReader : public Reader { } size_t tell() override { +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } // Warning: do not use this function from multiple threads as // it advances the file descriptor void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } // Warning: do not use this function from multiple threads as @@ -103,34 +117,49 @@ class ParallelFileReader : public Reader { } private: + // Reads larger than this are split in batches and read in parallel. static constexpr size_t batch_size_ = 1 << 25; - static ThreadPool& thread_pool(); + + // The pool that reads the batches, held from the first batched read until + // the reader is destroyed. It cannot be io::thread_pool(), the tasks that + // run there wait for these batches and would deadlock in the same pool. + ThreadPool& thread_pool(); + int fd_; std::string label_; + std::once_flag pool_once_; + std::shared_ptr pool_; }; class FileWriter : public Writer { public: explicit FileWriter() {} explicit FileWriter(std::string file_path) - : fd_(open( - file_path.c_str(), - O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, - 0644)), - label_(std::move(file_path)) {} + : file_path_(std::move(file_path)) {} FileWriter(const FileWriter&) = delete; FileWriter& operator=(const FileWriter&) = delete; - FileWriter(FileWriter&& other) { - std::swap(fd_, other.fd_); + FileWriter(FileWriter&& other) + : fd_(std::exchange(other.fd_, -1)), + file_path_(std::move(other.file_path_)) { + other.file_path_.clear(); } ~FileWriter() override { - if (fd_ != 0) { + if (fd_ >= 0) { close(fd_); } } + // Kept separate from construction so lazy inputs can be evaluated first, + // they may still read from the file. + void open() override { + if (fd_ < 0 && !file_path_.empty()) { + fd_ = ::open( + file_path_.c_str(), O_CREAT | O_WRONLY | O_TRUNC | O_BINARY, 0644); + } + } + bool is_open() const override { return fd_ >= 0; } @@ -140,21 +169,34 @@ class FileWriter : public Writer { } size_t tell() override { + check_open(); +#ifdef _WIN32 + return _lseeki64(fd_, 0, SEEK_CUR); +#else return lseek(fd_, 0, SEEK_CUR); +#endif } void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + check_open(); + int origin; if (way == std::ios_base::beg) { - lseek(fd_, off, SEEK_SET); + origin = SEEK_SET; } else if (way == std::ios_base::end) { - lseek(fd_, off, SEEK_END); + origin = SEEK_END; } else { - lseek(fd_, off, SEEK_CUR); + origin = SEEK_CUR; } +#ifdef _WIN32 + _lseeki64(fd_, off, origin); +#else + lseek(fd_, off, origin); +#endif } void write(const char* data, size_t n) override { + check_open(); while (n != 0) { auto m = ::write(fd_, data, std::min(n, static_cast(INT32_MAX))); if (m <= 0) { @@ -168,12 +210,18 @@ class FileWriter : public Writer { } std::string label() const override { - return "file " + label_; + return "file " + file_path_; } private: - int fd_{0}; - std::string label_; + void check_open() const { + if (!is_open()) { + throw std::runtime_error("[write] File " + file_path_ + " is not open."); + } + } + + int fd_{-1}; + std::string file_path_; }; } // namespace io diff --git a/Source/Cmlx/include-framework/mlx-memory.h b/Source/Cmlx/include-framework/mlx-memory.h index bfa2f22d..448aa8e6 100644 --- a/Source/Cmlx/include-framework/mlx-memory.h +++ b/Source/Cmlx/include-framework/mlx-memory.h @@ -1,11 +1,13 @@ #ifdef __cplusplus -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #pragma once #include +#include #include +#include namespace mlx::core { @@ -16,6 +18,14 @@ namespace mlx::core { * */ MLX_API size_t get_active_memory(); +/* Get the size of the buffers backing the given arrays in bytes. + * + * Each unique buffer is counted once. The full allocator size of each buffer + * is used, which can exceed the logical size of its arrays. The arrays must be + * evaluated before calling this function. + * */ +MLX_API size_t get_array_buffer_size(const std::vector& arrays); + /* Get the peak amount of used memory in bytes. * * The maximum memory used recorded from the beginning of the program @@ -43,6 +53,9 @@ MLX_API size_t get_cache_memory(); * When Metal is available the memory limit defaults to 1.5 times the maximum * recommended working set size reported by the device. * + * On native Windows with CUDA the memory limit is also bounded by the current + * video memory budget reported by the operating system. + * * Returns the previous memory limit. * */ MLX_API size_t set_memory_limit(size_t limit); diff --git a/Source/Cmlx/include-framework/mlx-ops.h b/Source/Cmlx/include-framework/mlx-ops.h index 96e3ecfe..5e3d6cc1 100644 --- a/Source/Cmlx/include-framework/mlx-ops.h +++ b/Source/Cmlx/include-framework/mlx-ops.h @@ -382,6 +382,9 @@ inline array transpose( return transpose(a, std::vector(axes), s); } +/** Permutes last two dimensions of an array. */ +MLX_API array matrix_transpose(const array& a, StreamOrDevice s = {}); + /** Swap two axes of an array. */ MLX_API array swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {}); @@ -784,6 +787,23 @@ min(const array& a, MLX_API array min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); +MLX_API std::tuple +minmax(const array& a, bool keepdims, StreamOrDevice s = {}); +inline std::tuple minmax(const array& a, StreamOrDevice s = {}) { + return minmax(a, false, to_stream(s)); +} + +/** The minimum and maximum of the elements of an array along the given axes. */ +MLX_API std::tuple minmax( + const array& a, + const std::vector& axes, + bool keepdims = false, + StreamOrDevice s = {}); + +/** The minimum and maximum of the elements of an array along the given axis. */ +MLX_API std::tuple +minmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); + /** Returns the Hanning window of size M. */ MLX_API array hanning(int M, StreamOrDevice s = {}); @@ -879,6 +899,27 @@ MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); /** Returns topk elements of the array along a given axis. */ MLX_API array topk(const array& a, int k, int axis, StreamOrDevice s = {}); +/** + * Returns the sorted unique elements of the flattened array, and optionally + * the index of the first occurrence of each, the inverse indices and the + * counts. The output has the given size, and + * is truncated if ``size`` is smaller than the number of unique values. + * If ``size`` is larger, the result is padded with ``fill_value``, or the + * first of the sorted unique values if no fill value is provided. An empty + * input has none, so it throws unless ``size`` is zero or a fill value is + * given. + * A truncated output loses the counts of the values it drops, and its inverse + * indices are clamped to the last entry. + */ +MLX_API std::vector unique( + const array& a, + int size, + bool return_index = false, + bool return_inverse = false, + bool return_counts = false, + const std::optional& fill_value = std::nullopt, + StreamOrDevice s = {}); + /** Cumulative logsumexp of an array. */ MLX_API array logcumsumexp( const array& a, @@ -1627,6 +1668,7 @@ MLX_API array gather_qmm( std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "affine", + const std::optional& global_scale = std::nullopt, bool sorted_indices = false, StreamOrDevice s = {}); diff --git a/Source/Cmlx/include-framework/mlx-primitives.h b/Source/Cmlx/include-framework/mlx-primitives.h index 33d2bb09..549b4096 100644 --- a/Source/Cmlx/include-framework/mlx-primitives.h +++ b/Source/Cmlx/include-framework/mlx-primitives.h @@ -1080,8 +1080,13 @@ class FFT : public UnaryPrimitive { Stream stream, const std::vector& axes, bool inverse, - bool real) - : UnaryPrimitive(stream), axes_(axes), inverse_(inverse), real_(real) {} + bool real, + bool odd_out = false) + : UnaryPrimitive(stream), + axes_(axes), + inverse_(inverse), + real_(real), + odd_out_(odd_out) {} void eval_cpu(const std::vector& inputs, array& out) override; void eval_gpu(const std::vector& inputs, array& out) override; @@ -1092,13 +1097,16 @@ class FFT : public UnaryPrimitive { bool is_equivalent(const Primitive& other) const override; auto state() const { - return std::make_tuple(axes_, inverse_, real_); + return std::make_tuple(axes_, inverse_, real_, odd_out_); } private: std::vector axes_; bool inverse_; bool real_; + // The output length of an inverse real transform is not recoverable from + // the input shape. + bool odd_out_; }; class Flatten : public UnaryPrimitive { @@ -1140,6 +1148,20 @@ class Floor : public UnaryPrimitive { DEFINE_INPUT_OUTPUT_SHAPE() }; +class FloorDivide : public UnaryPrimitive { + public: + explicit FloorDivide(Stream stream) : UnaryPrimitive(stream) {} + + void eval_cpu(const std::vector& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_VMAP() + DEFINE_GRADS() + DEFINE_NAME(FloorDivide) + DEFINE_DEFAULT_IS_EQUIVALENT() + DEFINE_INPUT_OUTPUT_SHAPE() +}; + class Full : public UnaryPrimitive { public: explicit Full(Stream stream) : UnaryPrimitive(stream) {} @@ -1889,6 +1911,7 @@ class Scan : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS(); + DEFINE_INPUT_OUTPUT_SHAPE() const char* name() const override { switch (reduce_type_) { diff --git a/Source/Cmlx/include-framework/mlx-stream.h b/Source/Cmlx/include-framework/mlx-stream.h index 9b9bae94..2d219bc3 100644 --- a/Source/Cmlx/include-framework/mlx-stream.h +++ b/Source/Cmlx/include-framework/mlx-stream.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -30,6 +31,9 @@ struct MLX_API ThreadLocalStream : public Stream { /** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); +/** Get the default stream of current thread if it exists. */ +MLX_API std::optional peek_default_stream(Device d); + /** Make the stream the default for its device on current thread. */ MLX_API void set_default_stream(Stream s); diff --git a/Source/Cmlx/include-framework/mlx-types-complex.h b/Source/Cmlx/include-framework/mlx-types-complex.h index c5c1f284..76466791 100644 --- a/Source/Cmlx/include-framework/mlx-types-complex.h +++ b/Source/Cmlx/include-framework/mlx-types-complex.h @@ -46,6 +46,11 @@ struct complex64_t : public std::complex { operator float() const { return real(); }; + + // Explicit, so int casts still go through operator float + explicit operator bool() const { + return real() != 0 || imag() != 0; + }; }; inline bool operator>=(const complex64_t& a, const complex64_t& b) { @@ -107,9 +112,24 @@ inline complex64_t operator-(const complex64_t& v) { complex_binop_helper(_op_, _operator_, float16_t) \ complex_binop_helper(_op_, _operator_, bfloat16_t) \ complex_binop_helper(_op_, _operator_, float) -// clang-format on complex_binop(+, operator+) +template +constexpr bool is_complex = false; + +template +constexpr bool is_complex().real())>> = true; +// clang-format on + +template +inline bool isnan(T v) { + if constexpr (is_complex) { + return std::isnan(std::real(v)) || std::isnan(std::imag(v)); + } else { + return std::isnan(v); + } +} + } // namespace mlx::core #endif diff --git a/Source/Cmlx/include-framework/mlx-types-half_types.h b/Source/Cmlx/include-framework/mlx-types-half_types.h index b6869f52..49140ca1 100644 --- a/Source/Cmlx/include-framework/mlx-types-half_types.h +++ b/Source/Cmlx/include-framework/mlx-types-half_types.h @@ -57,13 +57,6 @@ fp16_bf16_binop_helper(/, operator/) #endif -template -constexpr bool is_complex = false; - -template -constexpr bool is_complex().real())>> = - true; - template constexpr bool is_signed_v = std::is_signed_v || std::is_same_v || std::is_same_v; diff --git a/Source/Cmlx/include-framework/mlx-utils.h b/Source/Cmlx/include-framework/mlx-utils.h index a0d8478a..53665bb8 100644 --- a/Source/Cmlx/include-framework/mlx-utils.h +++ b/Source/Cmlx/include-framework/mlx-utils.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -29,7 +30,9 @@ MLX_API Stream to_stream(StreamOrDevice s, Device default_); struct StreamContext { public: - StreamContext(StreamOrDevice s) : _stream(default_stream(default_device())) { + StreamContext(StreamOrDevice s) + : _stream(default_stream(default_device())), + _owner(std::this_thread::get_id()) { if (std::holds_alternative(s)) { throw std::runtime_error( "[StreamContext] Invalid argument, please specify a stream or device."); @@ -39,13 +42,21 @@ struct StreamContext { set_default_stream(_s); } - ~StreamContext() { + // noexcept(false) so the thread check below can throw. A destructor that + // throws while another exception unwinds this thread calls std::terminate. + ~StreamContext() noexcept(false) { + if (std::this_thread::get_id() != _owner) { + throw std::runtime_error( + "[StreamContext] Destroyed in a different thread than it was " + "constructed in."); + } set_default_device(_stream.device); set_default_stream(_stream); } private: Stream _stream; + std::thread::id _owner; }; struct MLX_API PrintOptions { diff --git a/Source/Cmlx/include-framework/mlx-version.h b/Source/Cmlx/include-framework/mlx-version.h index 3b3eb3f2..f0275828 100644 --- a/Source/Cmlx/include-framework/mlx-version.h +++ b/Source/Cmlx/include-framework/mlx-version.h @@ -7,7 +7,7 @@ #define MLX_VERSION_MAJOR 0 #define MLX_VERSION_MINOR 32 -#define MLX_VERSION_PATCH 2 +#define MLX_VERSION_PATCH 4 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) diff --git a/Source/Cmlx/include/mlx/c/distributed_group.h b/Source/Cmlx/include/mlx/c/distributed_group.h index bfbaa80e..a784cf63 100644 --- a/Source/Cmlx/include/mlx/c/distributed_group.h +++ b/Source/Cmlx/include/mlx/c/distributed_group.h @@ -4,6 +4,7 @@ #define MLX_DISTRIBUTED_GROUP_H #include +#include #include "mlx/c/stream.h" @@ -41,15 +42,67 @@ int mlx_distributed_init( bool strict, const char* bk /* may be null */); +/** + * Byte-level all-gather callback used by the JACCL side channel. + * + * \param rank The rank of the process. + * \param size The size of the group. + * \param src Pointer to this rank's data of size n_bytes. + * \param dst Pointer to an output buffer of size size * n_bytes. After the + * call, dst[rank * n_bytes, (rank+1) * n_bytes] contains the data from rank. + * \param n_bytes The number of bytes contributed by each rank. + * \param payload Opaque user data passed to the callback. + * \return Zero on success, non-zero on failure. + */ +typedef int (*mlx_distributed_all_gather_fn)( + int rank, + int size, + const char* src, + char* dst, + size_t n_bytes, + void* payload); + +/** + * Initialize JACCL with a custom metadata all-gather callback. + * + * `bk` must be `jaccl`. The synchronous callback reads `n_bytes` from `src` + * and fills `size * n_bytes` in `dst`, ordered by `rank`. Buffers are borrowed. + * The callback returns zero on success, nonzero on failure. + * + * If a `jaccl` group is already cached, it is returned and `all_gather` is + * never called: the callback is dropped and, with `dtor`, `payload` is + * destroyed before this function returns. Call + * `mlx_distributed_clear_backends()` first to initialize with a new callback. + * + * With `dtor`, ownership of `payload` transfers on entry, including failure. + * Destruction occurs once after all references are released. Cached or + * live groups may delay it. Without `dtor`, `payload` is borrowed and must + * remain valid while the callback is retained. + * + * Returns zero on success, nonzero on error. `res` is unchanged on failure. + */ +int mlx_distributed_init_with_all_gather( + mlx_distributed_group* res, + bool strict, + const char* bk, + mlx_distributed_all_gather_fn all_gather, + void* payload, + void (*dtor)(void*)); + +/** + * Clear the distributed backend cache. + */ +int mlx_distributed_clear_backends(void); + /** * Get the rank. */ -int mlx_distributed_group_rank(mlx_distributed_group group); +int mlx_distributed_group_get_rank(int* res, mlx_distributed_group group); /** * Get the group size. */ -int mlx_distributed_group_size(mlx_distributed_group group); +int mlx_distributed_group_get_size(int* res, mlx_distributed_group group); /** * Split the group. diff --git a/Source/Cmlx/include/mlx/c/fast.h b/Source/Cmlx/include/mlx/c/fast.h index 914fd2a2..e4445f72 100644 --- a/Source/Cmlx/include/mlx/c/fast.h +++ b/Source/Cmlx/include/mlx/c/fast.h @@ -96,6 +96,16 @@ int mlx_fast_cuda_kernel_apply( const mlx_fast_cuda_kernel_config config, const mlx_stream stream); +int mlx_fast_gated_delta_update( + mlx_vector_array* res, + const mlx_array queries, + const mlx_array keys, + const mlx_array values, + const mlx_array gates, + const mlx_array beta_, + const mlx_array initial_state /* may be null */, + const mlx_array mask /* may be null */, + const mlx_stream s); int mlx_fast_layer_norm( mlx_array* res, const mlx_array x, diff --git a/Source/Cmlx/include/mlx/c/memory.h b/Source/Cmlx/include/mlx/c/memory.h index bae9e08e..972142c7 100644 --- a/Source/Cmlx/include/mlx/c/memory.h +++ b/Source/Cmlx/include/mlx/c/memory.h @@ -30,6 +30,7 @@ extern "C" { int mlx_clear_cache(void); int mlx_get_active_memory(size_t* res); +int mlx_get_array_buffer_size(size_t* res, const mlx_vector_array arrays); int mlx_get_cache_memory(size_t* res); int mlx_get_memory_limit(size_t* res); int mlx_get_peak_memory(size_t* res); diff --git a/Source/Cmlx/include/mlx/c/ops.h b/Source/Cmlx/include/mlx/c/ops.h index 74177857..23359f63 100644 --- a/Source/Cmlx/include/mlx/c/ops.h +++ b/Source/Cmlx/include/mlx/c/ops.h @@ -548,6 +548,7 @@ int mlx_gather_qmm( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, bool sorted_indices, const mlx_stream s); int mlx_gather_qqmm( @@ -703,6 +704,7 @@ int mlx_matmul( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_matrix_transpose(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_max_axes( mlx_array* res, const mlx_array a, @@ -1360,6 +1362,15 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unique( + mlx_vector_array* res, + const mlx_array a, + int size, + bool return_index, + bool return_inverse, + bool return_counts, + const mlx_array fill_value /* may be null */, + const mlx_stream s); int mlx_unstack_axis( mlx_vector_array* res, const mlx_array a, diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index 1f8e74e3..36969c09 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit 1f8e74e3f12f31365464a6867c6579f0e9b29d85 +Subproject commit 36969c094354086e15d2a327a9c50feb20b9cf11 diff --git a/Source/Cmlx/mlx-c b/Source/Cmlx/mlx-c index ebc88f10..19d3c9c1 160000 --- a/Source/Cmlx/mlx-c +++ b/Source/Cmlx/mlx-c @@ -1 +1 @@ -Subproject commit ebc88f10caa1b625e6b581437a8dea6df8a70085 +Subproject commit 19d3c9c180dddf51be441b55e7f454861899bc93 diff --git a/Source/Cmlx/mlx-generated/binary_ops.cpp b/Source/Cmlx/mlx-generated/binary_ops.cpp index 2c815772..9c2d6f47 100644 --- a/Source/Cmlx/mlx-generated/binary_ops.cpp +++ b/Source/Cmlx/mlx-generated/binary_ops.cpp @@ -273,7 +273,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/compiled_preamble.cpp b/Source/Cmlx/mlx-generated/compiled_preamble.cpp index 8b89f66a..48006a42 100644 --- a/Source/Cmlx/mlx-generated/compiled_preamble.cpp +++ b/Source/Cmlx/mlx-generated/compiled_preamble.cpp @@ -76,11 +76,6 @@ inline float operator+(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator-(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator*(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator*(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator/(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) / static_cast(rhs); } inline float operator/(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) / static_cast(rhs); } -template -constexpr bool is_complex = false; -template -constexpr bool is_complex().real())>> = - true; template constexpr bool is_signed_v = std::is_signed_v || std::is_same_v || std::is_same_v; @@ -120,6 +115,9 @@ struct complex64_t : public std::complex { operator float() const { return real(); }; + explicit operator bool() const { + return real() != 0 || imag() != 0; + }; }; inline bool operator>=(const complex64_t& a, const complex64_t& b) { return (a.real() > b.real()) || @@ -147,6 +145,18 @@ inline complex64_t operator-(const complex64_t& v) { return -static_cast>(v); } inline complex64_t operator+(const std::complex& x, const complex64_t& y) { return x + static_cast>(y); } inline complex64_t operator+(const complex64_t& x, const std::complex& y) { return static_cast>(x) + y; } inline complex64_t operator+(const complex64_t& x, const complex64_t& y) { return static_cast>(x) + static_cast>(y); } inline complex64_t operator+(bool x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, bool y) { return x + static_cast(y); } inline complex64_t operator+(uint32_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, uint32_t y) { return x + static_cast(y); } inline complex64_t operator+(uint64_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, uint64_t y) { return x + static_cast(y); } inline complex64_t operator+(int32_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, int32_t y) { return x + static_cast(y); } inline complex64_t operator+(int64_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, int64_t y) { return x + static_cast(y); } inline complex64_t operator+(float16_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, float16_t y) { return x + static_cast(y); } inline complex64_t operator+(bfloat16_t x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, bfloat16_t y) { return x + static_cast(y); } inline complex64_t operator+(float x, const complex64_t& y) { return static_cast(x) + y; } inline complex64_t operator+(const complex64_t& x, float y) { return x + static_cast(y); } +template +constexpr bool is_complex = false; +template +constexpr bool is_complex().real())>> = true; +template +inline bool isnan(T v) { + if constexpr (is_complex) { + return std::isnan(std::real(v)) || std::isnan(std::imag(v)); + } else { + return std::isnan(v); + } +} } namespace mlx::core::simd { template @@ -269,7 +279,7 @@ auto imag(Simd in) -> Simd { } template Simd isnan(Simd in) { - return std::isnan(in.value); + return mlx::core::isnan(in.value); } template auto operator +(Simd a, Simd b) ->Simd { return a.value + b.value; } template auto operator +(T1 a, Simd b)->Simd { return a + b.value; } template auto operator +(Simd a, T2 b)->Simd { return a.value + b; } template auto operator -(Simd a, Simd b) ->Simd { return a.value - b.value; } template auto operator -(T1 a, Simd b)->Simd { return a - b.value; } template auto operator -(Simd a, T2 b)->Simd { return a.value - b; } @@ -287,11 +297,46 @@ Simd clz(Simd x_) { return __builtin_clz(x_.value); } template +Simd divide(Simd a_, Simd b_) { + T a = a_.value; + T b = b_.value; + if constexpr (std::is_integral_v) { + if (b == 0) { + return T(0); + } + } + return a / b; +} +template +Simd floor_divide(Simd a_, Simd b_) { + T a = a_.value; + T b = b_.value; + if constexpr (std::is_integral_v) { + if (b == 0) { + return T(0); + } + T q = a / b; + if constexpr (is_signed_v) { + if (a % b != 0 && (a < 0 != b < 0)) { + q -= 1; + } + } + return q; + } else if constexpr (is_complex) { + return a / b; + } else { + return floor(divide(a_, b_)); + } +} +template Simd remainder(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; T r; if constexpr (std::is_integral_v) { + if (b == 0) { + return a; + } r = a % b; } else { r = std::remainder(a, b); @@ -308,7 +353,7 @@ Simd maximum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -319,7 +364,7 @@ Simd minimum(Simd a_, Simd b_) { T a = a_.value; T b = b_.value; if constexpr (!std::is_integral_v) { - if (std::isnan(a)) { + if (mlx::core::isnan(a)) { return a; } } @@ -640,19 +685,20 @@ struct ToFP8 { struct FromFP8 { template Simd operator()(Simd x) { - auto v = Simd(x & 127) << 7; + auto v = Simd(x & 127); + auto sign_bit = Simd((x >> 7) & 1) << 15; + auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; Simd out; if constexpr (simd::max_size >= N) { - auto converted = *(Simd*)(&v); + auto converted = *(Simd*)(&u); out = converted * 256.0; } else { for (int i = 0; i < N; ++i) { - auto converted = *(float16_t*)(&v[i]); + auto converted = *(float16_t*)(&u[i]); out[i] = converted * 256.0; } } - auto sign = Simd(x & 128); - return select(sign, -out, out); + return out; } float operator()(uint8_t x) { return (*this)(Simd(x)).value; @@ -663,7 +709,8 @@ namespace mlx::core::detail { using namespace mlx::core::simd; struct Add { template Simd operator()(Simd x, Simd y) { return operator+(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct ArcTan2 { template Simd operator()(Simd x, Simd y) { return atan2(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; -struct Divide { template Simd operator()(Simd x, Simd y) { return operator/(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; +struct Divide { template Simd operator()(Simd x, Simd y) { return divide(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; +struct FloorDivide { template Simd operator()(Simd x, Simd y) { return floor_divide(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct Multiply { template Simd operator()(Simd x, Simd y) { return operator*(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct Subtract { template Simd operator()(Simd x, Simd y) { return operator-(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; struct LogicalAnd { template Simd operator()(Simd x, Simd y) { return operator&&(x, y); } template T operator()(T x, T y) { return (*this)(Simd(x), Simd(y)).value; } }; diff --git a/Source/Cmlx/mlx-generated/conv.cpp b/Source/Cmlx/mlx-generated/conv.cpp index 25c6d543..3cdec867 100644 --- a/Source/Cmlx/mlx-generated/conv.cpp +++ b/Source/Cmlx/mlx-generated/conv.cpp @@ -224,8 +224,8 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -374,8 +374,8 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -561,6 +561,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -582,6 +583,7 @@ struct Conv2DWeightBlockLoader { weight_hw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -596,7 +598,7 @@ struct Conv2DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -658,9 +660,9 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; short kdil_d; short kdil_h; @@ -842,9 +844,9 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -1064,6 +1066,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -1085,6 +1088,7 @@ struct Conv3DWeightBlockLoader { weight_dhw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -1099,7 +1103,7 @@ struct Conv3DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -1206,6 +1210,7 @@ struct Conv2DInputBlockLoaderSmallChannels { // Rows / strided reads within the block STEEL_CONST short n_rows = BROWS / TROWS; + static_assert(n_rows > 0, "BM must be at least tgp_size * vec_size / BK"); // Thread location indices const short thread_idx; @@ -1368,6 +1373,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -1388,6 +1394,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { params(params_), weight_hw(thread_idx % TCOLS), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -1395,7 +1402,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { if (bi >= BROWS || bj >= BCOLS) return; - if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -1424,7 +1431,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } } else { for (short i = 0; i < BROWS; i += TROWS) { - if (((read_n + i) < params->O)) { + if (((read_n + i) < N)) { STEEL_PRAGMA_UNROLL for (short j = 0; j < n_channels; j++) { dst[i * dst_ld + j] = curr_src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/cross_entropy.cpp b/Source/Cmlx/mlx-generated/cross_entropy.cpp new file mode 100644 index 00000000..53d04699 --- /dev/null +++ b/Source/Cmlx/mlx-generated/cross_entropy.cpp @@ -0,0 +1,152 @@ +namespace mlx::core::metal { + +const char* cross_entropy() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/cross_entropy.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/cross_entropy.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/cross_entropy.h" +// Copyright © 2026 Apple Inc. + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy( + const device T* in, + const device int32_t* targets, + device float* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + in += gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(in[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = + (offset + i < axis_size) ? AccT(in[offset + i]) : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + uint n_simdgroups = ceildiv(lsize, uint(SIMD_SIZE)); + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = (simd_lane_id < n_simdgroups) ? local_max[simd_lane_id] + : Limits::finite_min; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = + (simd_lane_id < n_simdgroups) ? local_normalizer[simd_lane_id] : AccT(0); + normalizer = simd_sum(normalizer); + + if (lid == 0) { + // Subtract the two logits first; they are the same magnitude, so less is + // lost than in logsumexp - x_t. + AccT gap = maxval - x_t; + out[gid] = isinf(maxval) ? float(gap) : float(log(normalizer) + gap); + } +} + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy_vjp( + const device T* in, + const device int32_t* targets, + const device float* loss, + const device float* cotan, + device T* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + size_t row_offset = gid * size_t(axis_size); + in += row_offset; + out += row_offset; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + AccT g = AccT(cotan[gid]); + AccT l = AccT(loss[gid]); + + // out aliases in when the input is donated, so every thread must read the + // target column before any thread writes. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } else { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + if (col < axis_size) { + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } + } + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h index fe469a89..54d2d1e0 100644 --- a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -755,7 +755,7 @@ constexpr char jit_source_cast_op[] = { 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, - 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, 0x69, + 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x7c, 0x7c, 0x20, 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, @@ -772,7 +772,7 @@ constexpr char jit_source_cast_op[] = { 0x72, 0x28, 0x29, 0x28, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, - 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x31, 0x2c, 0x20, 0x31, 0x7d, 0x20, + 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x31, 0x2c, 0x20, 0x30, 0x7d, 0x20, 0x3a, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x30, 0x2c, 0x20, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6e, @@ -4466,143 +4466,155 @@ constexpr char jit_source_qmm_naive[] = { 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, - 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, - 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, - 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, - 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, - 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, - 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, - 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, - 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, - 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, - 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, - 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, - 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, - 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, - 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, - 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, - 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, 0x20, 0x20, - 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x61, - 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, 0x20, 0x62, - 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, - 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, 0x62, 0x6c, - 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, - 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, 0x3d, 0x20, - 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, 0x3b, 0x0a, - 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, - 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, - 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, 0x6e, 0x2c, - 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, - 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, - 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, - 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, - 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, - 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, - 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, - 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, - 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, - 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, - 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, - 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, - 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, - 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, - 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, - 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, - 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, - 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, - 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, - 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, - 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, - 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, - 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, - 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, 0x63, 0x74, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, 0x65, 0x20, 0x67, 0x6c, + 0x6f, 0x62, 0x61, 0x6c, 0x20, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x69, + 0x73, 0x20, 0x70, 0x65, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, + 0x65, 0x64, 0x20, 0x65, 0x78, 0x70, 0x65, 0x72, 0x74, 0x3b, 0x20, 0x61, + 0x20, 0x6e, 0x6f, 0x6e, 0x2d, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x65, + 0x64, 0x20, 0x63, 0x61, 0x6c, 0x6c, 0x20, 0x68, 0x61, 0x73, 0x20, 0x6f, + 0x6e, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x67, 0x6c, + 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x26, + 0x26, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, 0x6c, 0x6f, + 0x62, 0x61, 0x6c, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x2b, 0x3d, + 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, + 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x2e, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, + 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, + 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, + 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4e, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, + 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, + 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, + 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, + 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, 0x20, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x74, + 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, 0x3d, + 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, 0x6e, + 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, + 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x6c, + 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x42, + 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, + 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, + 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, + 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, + 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, + 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x5b, + 0x26, 0x5d, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, + 0x69, 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, + 0x70, 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, + 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, + 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, + 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, - 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x5b, 0x26, - 0x5d, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, - 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, - 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, - 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, - 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, - 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, - 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, - 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, - 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, - 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, - 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, - 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, - 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, - 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, - 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, - 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, - 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x75, 0x6d, 0x6d, 0x79, - 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3b, 0x20, 0x6e, 0x6f, 0x2d, - 0x62, 0x69, 0x61, 0x73, 0x20, 0x70, 0x61, 0x74, 0x68, 0x73, 0x20, 0x6e, - 0x65, 0x76, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, - 0x6f, 0x72, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x67, 0x5a, 0x2e, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, - 0x20, 0x67, 0x53, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, - 0x20, 0x7d, 0x28, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, - 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, 0x65, - 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, 0x6f, - 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, - 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x6d, - 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, - 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, - 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, 0x2a, 0x20, - 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, - 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x20, 0x2a, 0x20, - 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, - 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, - 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, - 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, - 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, 0x20, 0x2d, 0x20, 0x73, - 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, 0x2a, - 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x67, 0x41, 0x29, - 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, 0x69, - 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x3c, - 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x48, 0x61, 0x73, 0x4b, - 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x53, 0x4d, 0x38, - 0x30, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, - 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, - 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, - 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, - 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, - 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, - 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, - 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x75, 0x6d, 0x6d, + 0x79, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3b, 0x20, 0x6e, 0x6f, + 0x2d, 0x62, 0x69, 0x61, 0x73, 0x20, 0x70, 0x61, 0x74, 0x68, 0x73, 0x20, + 0x6e, 0x65, 0x76, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x6f, 0x72, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x67, 0x5a, 0x2e, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x67, 0x53, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x28, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, + 0x6f, 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, + 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, + 0x6d, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, 0x2a, + 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x20, 0x2a, + 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, 0x20, 0x42, 0x4c, + 0x4b, 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, + 0x73, 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, 0x20, 0x2d, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, + 0x2a, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x67, 0x41, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, + 0x69, 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, + 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x48, 0x61, 0x73, + 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x53, 0x4d, + 0x38, 0x30, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, + 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_qmm_sm80[] = { @@ -5482,47 +5494,288 @@ constexpr char jit_source_qmm_sm80[] = { 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x20, 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x28, 0x74, 0x43, 0x72, 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x28, 0x69, 0x29, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, - 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2c, 0x20, - 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, 0x2c, 0x20, 0x72, 0x32, - 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x5f, - 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, - 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, - 0x66, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, - 0x2c, 0x20, 0x74, 0x43, 0x70, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, - 0x74, 0x43, 0x73, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, - 0x67, 0x43, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, - 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, - 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, - 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, - 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, - 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3d, 0x20, 0x49, 0x6e, 0x74, - 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x7b, - 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6d, 0x61, 0x6b, - 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, - 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, - 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, - 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, - 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x30, - 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, + 0x65, 0x20, 0x6c, 0x61, 0x73, 0x74, 0x20, 0x70, 0x72, 0x65, 0x66, 0x65, + 0x74, 0x63, 0x68, 0x65, 0x73, 0x20, 0x6d, 0x61, 0x79, 0x20, 0x73, 0x74, + 0x69, 0x6c, 0x6c, 0x20, 0x62, 0x65, 0x20, 0x6c, 0x61, 0x6e, 0x64, 0x69, + 0x6e, 0x67, 0x20, 0x69, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x62, 0x75, + 0x66, 0x66, 0x65, 0x72, 0x73, 0x20, 0x43, 0x20, 0x72, 0x65, 0x75, 0x73, + 0x65, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x63, 0x70, 0x5f, 0x61, 0x73, 0x79, + 0x6e, 0x63, 0x5f, 0x77, 0x61, 0x69, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, + 0x6f, 0x70, 0x79, 0x28, 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x63, 0x2c, 0x20, 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, + 0x2c, 0x20, 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x63, 0x2c, 0x20, 0x74, 0x43, 0x70, 0x43, 0x2c, 0x20, + 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x2c, 0x20, 0x73, 0x32, + 0x67, 0x5f, 0x74, 0x43, 0x67, 0x43, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, + 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3d, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, + 0x7a, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, - 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, - 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, + 0x7a, 0x65, 0x29, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x28, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x3c, 0x5f, 0x30, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, + 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, + 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, + 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, 0x61, + 0x75, 0x6e, 0x63, 0x68, 0x5f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, 0x5f, + 0x5f, 0x28, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x29, 0x29, 0x29, 0x3a, 0x3a, + 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6b, 0x65, 0x72, + 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, + 0x6c, 0x65, 0x2a, 0x20, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2a, 0x20, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, + 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, + 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, + 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, 0x61, + 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, + 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, + 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x7a, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, + 0x69, 0x6e, 0x65, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, + 0x28, 0x6d, 0x69, 0x78, 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, + 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, + 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, + 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, + 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, + 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4e, 0x2c, + 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x49, 0x6e, + 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, + 0x6e, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, + 0x4e, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, + 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, + 0x6e, 0x64, 0x6c, 0x65, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, + 0x73, 0x74, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, + 0x32, 0x3e, 0x28, 0x64, 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x29, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, + 0x65, 0x73, 0x65, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, + 0x6c, 0x6c, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, + 0x6d, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, + 0x70, 0x74, 0x72, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, + 0x29, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, + 0x64, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, + 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, + 0x43, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, + 0x20, 0x6e, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, + 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, + 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, + 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x46, 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, + 0x72, 0x2c, 0x20, 0x75, 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x6c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x69, 0x6e, 0x70, 0x75, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, + 0x73, 0x6c, 0x69, 0x63, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, 0x73, 0x5f, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, + 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x20, 0x3f, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, + 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x61, + 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, 0x5f, 0x6e, + 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, + 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, + 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, + 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, + 0x20, 0x74, 0x68, 0x65, 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, + 0x69, 0x61, 0x74, 0x65, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x20, 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, + 0x65, 0x72, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, + 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x6d, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, + 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, + 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, + 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, + 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, + 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, + 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, + 0x6d, 0x53, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, + 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, + 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, + 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, + 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x20, 0x3d, 0x20, 0x6d, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x30, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, + 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, + 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, + 0x70, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x53, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, - 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, - 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, - 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x74, 0x79, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, 0x61, 0x75, 0x6e, 0x63, 0x68, @@ -5531,223 +5784,324 @@ constexpr char jit_source_qmm_sm80[] = { 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x29, 0x29, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, - 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, - 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x41, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x51, 0x75, 0x61, - 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2a, 0x20, - 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x5a, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, - 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x6c, 0x68, 0x73, - 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, - 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, - 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x43, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x2c, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, - 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, - 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, - 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, - 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, - 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, - 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, - 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, - 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, - 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x79, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x5f, 0x63, 0x6f, - 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, - 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x7a, 0x29, 0x3b, 0x0a, 0x0a, - 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, - 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x6d, 0x69, 0x78, - 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, - 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, - 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, - 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, - 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, - 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, - 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, - 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, - 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, - 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, - 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, - 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, - 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, - 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, - 0x75, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, - 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x3b, 0x0a, - 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, 0x6e, 0x64, 0x6c, 0x65, - 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x69, 0x6e, - 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x62, 0x72, 0x6f, - 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x64, - 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, - 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x29, - 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, 0x65, 0x73, 0x65, 0x6e, - 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, 0x6c, 0x6c, 0x20, 0x74, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, - 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x20, + 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x72, 0x68, 0x73, 0x5f, 0x6b, 0x65, 0x72, + 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, + 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, + 0x6c, 0x65, 0x2a, 0x20, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2a, 0x20, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, + 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2a, 0x20, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6e, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6b, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, + 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, + 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x20, 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, + 0x72, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x20, 0x3d, 0x20, 0x6d, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x2a, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, + 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x6d, 0x20, 0x2d, 0x20, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, + 0x20, 0x3d, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3c, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x20, + 0x3f, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3a, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x6d, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, + 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x6d, 0x69, + 0x78, 0x65, 0x64, 0x29, 0x2e, 0x20, 0x41, 0x6c, 0x6c, 0x20, 0x67, 0x61, + 0x74, 0x68, 0x65, 0x72, 0x65, 0x64, 0x20, 0x72, 0x6f, 0x77, 0x73, 0x20, + 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x6f, 0x6e, 0x65, 0x20, 0x6d, 0x61, 0x74, + 0x72, 0x69, 0x78, 0x3b, 0x20, 0x65, 0x61, 0x63, 0x68, 0x20, 0x72, 0x75, + 0x6e, 0x20, 0x6f, 0x66, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x65, 0x71, + 0x75, 0x61, 0x6c, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x20, 0x69, 0x6e, 0x73, 0x69, 0x64, 0x65, 0x20, 0x74, + 0x68, 0x65, 0x20, 0x72, 0x6f, 0x77, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, + 0x72, 0x65, 0x75, 0x73, 0x65, 0x73, 0x20, 0x6f, 0x6e, 0x65, 0x20, 0x65, + 0x78, 0x70, 0x65, 0x72, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, + 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, + 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, + 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, + 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, + 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, + 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x47, 0x72, + 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x28, 0x6e, 0x2c, 0x20, + 0x6b, 0x2c, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, + 0x6e, 0x74, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, + 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x31, 0x29, 0x2c, 0x20, 0x64, 0x41, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x3c, 0x51, 0x75, + 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, 0x29, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6b, + 0x2c, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x63, 0x6f, 0x75, 0x6e, + 0x74, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, - 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, - 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, - 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x3c, - 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, 0x29, 0x2c, 0x20, 0x6d, - 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, - 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, - 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, - 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, - 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, - 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, - 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, - 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, - 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, - 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, - 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, - 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, - 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, - 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, - 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, - 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, - 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, - 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, - 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x46, - 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x2c, 0x20, 0x75, - 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x6c, 0x6f, 0x6f, - 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x69, 0x6e, 0x70, 0x75, - 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, - 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, - 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, - 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, - 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, - 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, - 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, - 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, - 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, - 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, - 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, - 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, - 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, - 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, - 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, - 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, - 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, - 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, - 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, - 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, - 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, - 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, - 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, - 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, - 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, - 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, - 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x29, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, - 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, - 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, + 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, 0x31, 0x29, 0x2c, 0x20, + 0x64, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, + 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, + 0x28, 0x53, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, + 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x5a, + 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, - 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, - 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, - 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, - 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, - 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, - 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, - 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, - 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, - 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, - 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, - 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, - 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, - 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, - 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, - 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, - 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, - 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, - 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, - 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, - 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, + 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, + 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, + 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, + 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x77, 0x68, 0x69, 0x6c, 0x65, 0x20, 0x28, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3c, 0x20, + 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x3d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, + 0x78, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x5f, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x67, 0x70, + 0x5f, 0x62, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x6f, 0x77, 0x20, 0x3d, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x20, 0x31, 0x3b, 0x20, + 0x72, 0x6f, 0x77, 0x20, 0x3c, 0x20, 0x74, 0x67, 0x70, 0x5f, 0x62, 0x6d, + 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x6f, 0x77, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, + 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x72, 0x68, + 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x79, 0x5f, + 0x72, 0x6f, 0x77, 0x20, 0x2b, 0x20, 0x72, 0x6f, 0x77, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x21, 0x3d, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x20, 0x3d, 0x20, 0x72, 0x6f, 0x77, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5f, 0x6e, + 0x65, 0x78, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x72, 0x65, + 0x61, 0x6b, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, + 0x3d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x6e, 0x65, 0x78, + 0x74, 0x20, 0x2d, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x41, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, + 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x79, 0x5f, 0x72, 0x6f, + 0x77, 0x20, 0x2b, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, + 0x30, 0x29, 0x2c, 0x20, 0x6d, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x79, 0x5f, 0x72, 0x6f, 0x77, 0x20, 0x2b, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x30, 0x29, 0x2c, + 0x20, 0x6d, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, + 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, + 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, 0x6e, + 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, + 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, + 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, + 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x75, + 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x20, 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, - 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, - 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x41, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x72, 0x75, + 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, + 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, 0x20, 0x72, + 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, + 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, - 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, - 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, - 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, - 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, - 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, - 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, - 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, - 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, - 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, - 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, - 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, - 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, - 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, - 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, - 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, - 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x5a, 0x20, - 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, - 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, - 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, - 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, - 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, - 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, - 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, - 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, - 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, 0x6f, 0x72, - 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, - 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x5f, 0x6d, - 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, - 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x67, - 0x41, 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, - 0x5f, 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, - 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x38, 0x30, - 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x28, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, - 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x42, 0x2c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, - 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, - 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, - 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, - 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, + 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, + 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, + 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, + 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, + 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x5a, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x65, + 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, + 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, + 0x38, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x75, 0x6e, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x42, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, + 0x77, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x72, 0x6f, + 0x77, 0x73, 0x20, 0x3d, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, + 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x3a, 0x3a, 0x76, + 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3e, 0x20, + 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, + 0x20, 0x3c, 0x3d, 0x20, 0x31, 0x36, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, + 0x6d, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x2c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x28, + 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3c, 0x3d, 0x20, + 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x6e, + 0x74, 0x3c, 0x33, 0x32, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3e, 0x20, 0x31, 0x36, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, + 0x75, 0x6e, 0x5f, 0x72, 0x6f, 0x77, 0x73, 0x20, 0x3c, 0x3d, 0x20, 0x31, + 0x36, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x6e, 0x74, + 0x3c, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x75, 0x6e, 0x5f, 0x71, 0x6d, 0x6d, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x68, 0x65, + 0x20, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x72, 0x75, 0x6e, 0x27, 0x73, 0x20, + 0x6c, 0x6f, 0x61, 0x64, 0x73, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x6e, + 0x6f, 0x74, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x77, 0x72, 0x69, 0x74, 0x65, + 0x20, 0x43, 0x20, 0x77, 0x68, 0x69, 0x6c, 0x65, 0x20, 0x69, 0x74, 0x20, + 0x69, 0x73, 0x20, 0x73, 0x74, 0x69, 0x6c, 0x6c, 0x20, 0x62, 0x65, 0x69, + 0x6e, 0x67, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x2e, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_qmm_sm90[] = { diff --git a/Source/Cmlx/mlx-generated/fft.cpp b/Source/Cmlx/mlx-generated/fft.cpp index d60d565b..36c59094 100644 --- a/Source/Cmlx/mlx-generated/fft.cpp +++ b/Source/Cmlx/mlx-generated/fft.cpp @@ -664,10 +664,20 @@ constexpr complex_t operator*(complex_t a, complex_t b) { template constexpr complex_t operator/(complex_t a, complex_t b) { - auto denom = b.real * b.real + b.imag * b.imag; - auto x = a.real * b.real + a.imag * b.imag; - auto y = a.imag * b.real - a.real * b.imag; - return {x / denom, y / denom}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -675,11 +685,7 @@ template < typename U, enable_if_t && is_lane_convertible_v, bool> = true> constexpr complex_t operator/(U a, complex_t b) { - auto scalar = static_cast(a); - auto denom = b.real * b.real + b.imag * b.imag; - auto x = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx-generated/fp_quantized.cpp b/Source/Cmlx/mlx-generated/fp_quantized.cpp index 1c56cdb5..ad0c884b 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized.cpp @@ -98,11 +98,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -299,13 +299,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { } template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -317,7 +316,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -327,9 +327,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -343,6 +348,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -350,7 +358,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -365,14 +374,19 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales( scales_ + bi * src_ld / group_size + - (bj * pack_factor) / group_size) {} + (bj * pack_factor) / group_size) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -384,21 +398,15 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -489,6 +497,7 @@ template < typename T, int group_size, int bits, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( @@ -522,6 +531,15 @@ METAL_FUNC void fp_qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // Reuse the last valid row for reads beyond the output tile. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; @@ -531,8 +549,12 @@ METAL_FUNC void fp_qmv_fast_impl( load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device auto* sl = scales + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device auto* sl = scales + src * in_vec_size_g; U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); @@ -550,7 +572,7 @@ METAL_FUNC void fp_qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { if constexpr (has_global_scale) { y[row] = static_cast(result[row] * inv_scale_enc); } else { @@ -951,14 +973,17 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -983,7 +1008,7 @@ METAL_FUNC void fp_qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -994,7 +1019,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -1013,12 +1039,15 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -1028,7 +1057,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -1038,9 +1067,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -1050,7 +1088,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -1061,6 +1099,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -1076,12 +1123,14 @@ template < typename T, int group_size, int bits, + bool has_global_scale = false, int BM = 32, int BK = 32, int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -1119,7 +1168,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; auto wl = (const device uint8_t*)w; @@ -1134,16 +1184,23 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1152,9 +1209,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1172,7 +1239,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1181,9 +1252,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1199,8 +1280,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } @@ -1241,11 +1322,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device uint8_t*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -1288,6 +1370,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -1335,6 +1421,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( @@ -1376,6 +1463,7 @@ template < T, group_size, bits, + partial_rows, has_global_scale, results_per_simdgroup>( w, @@ -1395,6 +1483,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv( @@ -1603,16 +1692,18 @@ template template < typename T, - const int group_size, - const int bits, - const bool aligned_N, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool aligned_N, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -1652,8 +1743,31 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1709,8 +1823,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1739,10 +1853,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1759,7 +1874,7 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( + fp_qmv_fast_impl( w, scales, global_scale, @@ -1798,10 +1913,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1857,10 +1973,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1896,12 +2013,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> [[kernel]] void fp_gather_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1931,10 +2050,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1951,8 +2071,31 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1967,7 +2110,7 @@ template < const device uint32_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], const device T* x [[buffer(2)]], - device T* y [[buffer(3)]], + device float* y [[buffer(3)]], const constant int& K [[buffer(4)]], const constant int& N [[buffer(5)]], const constant int& M [[buffer(6)]], @@ -1992,9 +2135,10 @@ template < scales += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - fp_qmm_t_impl( + fp_qmm_t_impl( (const device uint32_t*)wl, scales, + nullptr, x, y, Xs, @@ -2020,6 +2164,7 @@ template < [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -2050,10 +2195,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -2070,8 +2216,21 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -2083,16 +2242,19 @@ template < int BK, int WM, int WN, - bool transpose> + bool transpose, + bool has_global_scale = false> [[kernel]] void fp_gather_qmm_rhs( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2123,7 +2285,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -2136,13 +2299,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2158,112 +2326,61 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp index 6b8473d8..78b04525 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp @@ -98,11 +98,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -163,7 +163,7 @@ inline float mx_scale_round_up(float x) { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/fp_quantized_nax.h" -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -223,13 +223,12 @@ struct Dequantize { }; template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -241,7 +240,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -269,6 +269,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -276,7 +279,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -289,7 +293,11 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id) {} + scales(scales_ + bi * src_ld / group_size + group_id) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -298,7 +306,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -312,14 +321,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -328,7 +330,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -354,6 +357,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -363,6 +367,7 @@ template < METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup Wtype* Ws, @@ -392,7 +397,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -408,7 +414,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -498,6 +504,7 @@ template < typename T, const int group_size, const int bits, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -507,6 +514,7 @@ template < METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Ws, @@ -536,7 +544,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -554,7 +563,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -649,11 +658,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device S*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -696,6 +706,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -710,10 +724,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -752,8 +768,19 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -811,8 +838,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -825,10 +852,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -857,10 +886,11 @@ template < threadgroup Wtype Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -877,8 +907,19 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -890,10 +931,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_n_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -924,10 +967,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -944,8 +988,8 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -958,16 +1002,19 @@ template < int WM, int WN, bool transpose, + bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_rhs_nax( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -984,7 +1031,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -996,13 +1044,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1028,10 +1081,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1039,151 +1092,125 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; - } - } + NAXTile Dtile; + Dtile.clear(); - xn += BK; - loader_w.next(); + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); } - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + xn += BK; + loader_w.next(); + } + + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp new file mode 100644 index 00000000..94a90fe8 --- /dev/null +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax.cpp @@ -0,0 +1,1968 @@ +namespace mlx::core::metal { + +const char* gated_delta_update_nax() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/gated_delta_update_nax.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/defines.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/defines.h" +// Copyright © 2024 Apple Inc. + + +#define STEEL_CONST static constant constexpr const +#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") +#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/type_traits.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/type_traits.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace metal { + +template +struct is_empty : metal::bool_constant<__is_empty(T)> {}; + +#ifdef __cpp_variable_templates +template +constexpr constant bool is_empty_v = is_empty::value; +#endif + +template +struct make_void { + typedef void type; +}; + +template +using void_t = typename make_void::type; + +template +struct is_static : metal::bool_constant>::value> {}; + +template +struct pointer_element {}; + +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; + +template +using pointer_element_t = typename pointer_element>::type; + +} // namespace metal + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// Integral constant with casting +/////////////////////////////////////////////////////////////////////////////// + +template +struct integral_constant { + static constexpr constant T value = v; + using value_type = T; + using type = integral_constant; + + METAL_FUNC constexpr operator value_type() const thread noexcept { + return value; + } +}; + +template +using bool_constant = integral_constant; +using true_type = bool_constant; +using false_type = bool_constant; + +template +struct is_integral : bool_constant::value> {}; + +template +struct is_integral> + : bool_constant::value> {}; + +template +constexpr constant bool is_integral_v = is_integral::value; + +template +using Int = integral_constant; + +/////////////////////////////////////////////////////////////////////////////// +// Binary Operators on Integral constants +/////////////////////////////////////////////////////////////////////////////// + +#define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ + } + +integral_const_binop(+, operator+); +integral_const_binop(-, operator-); +integral_const_binop(*, operator*); +integral_const_binop(/, operator/); + +integral_const_binop(==, operator==); +integral_const_binop(!=, operator!=); +integral_const_binop(<, operator<); +integral_const_binop(>, operator>); +integral_const_binop(<=, operator<=); +integral_const_binop(>=, operator>=); + +integral_const_binop(&&, operator&&); +integral_const_binop(||, operator||); + +template >> +METAL_FUNC constexpr auto operator||(true_type, T) { + return true_type{}; +} +template >> +METAL_FUNC constexpr auto operator||(T, true_type) { + return true_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(false_type, T) { + return false_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(T, false_type) { + return false_type{}; +} + +// Dispatch utilities +template +void dispatch_bool(bool v, F f) { + if (v) { + f(true_type{}); + } else { + f(false_type{}); + } +} + +template +constexpr void const_for_loop(F f) { + if constexpr (start < stop) { + constexpr auto idx = Int{}; + f(idx); + const_for_loop(f); + } +} + +#undef integral_const_binop + +/////////////////////////////////////////////////////////////////////////////// +// Reduction operators +/////////////////////////////////////////////////////////////////////////////// + +template +METAL_FUNC constexpr T sum(T x) { + return x; +} + +template +METAL_FUNC constexpr auto sum(T x, Us... us) { + return x + sum(us...); +} + +} // namespace steel +} // namespace mlx + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/nax.h" +// Copyright © 2025 Apple Inc. + + +#include +#include +#include + + +#include + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// MMA helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// NAX Steel with new tiles +/////////////////////////////////////////////////////////////////////////////// + +struct BaseNAXFrag { + STEEL_CONST short kFragRows = 16; + STEEL_CONST short kFragCols = 16; + + STEEL_CONST short kElemsPerFrag = (kFragRows * kFragCols) / 32; + + STEEL_CONST short kElemRows = 2; + STEEL_CONST short kElemCols = 4; + + STEEL_CONST short kElemRowsJump = 8; + + static_assert( + kElemRows * kElemCols == kElemsPerFrag, + "MMAFrag shape is not consistent with MMAFrag size"); + + template + using dtype_frag_t = typename metal::vec; + + METAL_FUNC static short2 get_coord() { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4; + return short2{fn, fm}; + } + + METAL_FUNC static short2 get_coord(short idx) { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)) + (idx >> 2) * 8; + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4 + idx % 4; + return short2{fn, fm}; + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + c + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_rows( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j)]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_safe( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((r < lx) && ((c + j) < ly)) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } else { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_rows( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_safe( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if (r < lx && (c + j) < ly) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename StartX, + typename StopX, + typename StartY, + typename StopY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_slice( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + StartX start_x, + StopX stop_x, + StartY start_y, + StopY stop_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + + const_for_loop<0, kElemRows, 1>([&](auto idx_row) { + const auto r = off_x + idx_row * Int{}; + if (r >= stop_x - sc.y || r < start_x - sc.y) { + return; + } + + const_for_loop<0, kElemCols, 1>([&](auto idx_col) { + const auto c = off_y + idx_col; + if (c >= stop_y - sc.x || c < start_y - sc.x) { + return; + } + + const auto src_idx = idx_row * Int{} + idx_col; + dst[(r + sc.y) * str_x + (c + sc.x) * str_y] = + static_cast(src[src_idx]); + }); + }); + } + + template + METAL_FUNC static constexpr void row_reduce( + thread const dtype_frag_t& inp_vals, + thread T* reduced_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + T thr_reduce = Op::apply( + Op::apply(inp_vals[i * kElemCols + 0], inp_vals[i * kElemCols + 1]), + Op::apply(inp_vals[i * kElemCols + 2], inp_vals[i * kElemCols + 3])); + + T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1)); + qgr_reduce = Op::apply(thr_reduce, qgr_reduce); + + T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8)); + sgr_reduce = Op::apply(qgr_reduce, sgr_reduce); + + reduced_vals[i] = Op::apply(reduced_vals[i], sgr_reduce); + } + } + + template + METAL_FUNC static constexpr void row_bin_op( + thread dtype_frag_t& inp_vals, + thread T* row_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + inp_vals[i * kElemCols + j] = + Op::apply(inp_vals[i * kElemCols + j], row_vals[i]); + } + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; + } + } +}; + +template < + typename T, + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> +struct NAXTile { + using NAXFrag_t = NAXFrag_; + using elem_type = T; + + STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; + STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; + STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; + + STEEL_CONST short kTileRows = kTileRows_; + STEEL_CONST short kTileCols = kTileCols_; + + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; + + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; + + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; + + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; + + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; + + METAL_FUNC NAXTile() thread {} + + METAL_FUNC constexpr void clear() thread { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kNumFrags; ++i) { + val_frags[i] = frag_type(0); + } + } + + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; + } + + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j) const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); + } + + template + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_reduce( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_bin_op( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_rows( + frag_at(), + src, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_safe( + frag_at(), + src, + ld, + Int<1>{}, + src_tile_dims.y, + src_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), + dst, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), + dst, + ld, + Int<1>{}, + dst_tile_dims.y, + dst_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_slice( + device U* dst, + const int ld, + const short2 start, + const short2 stop) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_slice( + frag_at(), + dst, + ld, + Int<1>{}, + start.y, + stop.y, + start.x, + stop.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } +}; + +template < + class CTile, + class ATile, + class BTile, + bool transpose_a, + bool transpose_b> +METAL_FUNC void tile_matmad_nax( + thread CTile& C, + thread ATile& A, + metal::bool_constant, + thread BTile& B, + metal::bool_constant) { + // Static checks + constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); + + constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); + + constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); + + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; + + if constexpr (TN == 1 && TM % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm += 2) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } + } + } + } +} + +} // namespace steel +} // namespace mlx + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +// Copyright © 2026 Apple Inc. + +#include + +#include +#include + + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_update_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax.h" +// Copyright © 2026 Apple Inc. + + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +/////////////////////////////////////////////////////////////////////////////// +// Function constants +/////////////////////////////////////////////////////////////////////////////// + +constant bool save_state_cache [[function_constant(200)]]; + +template +[[kernel]] void gated_delta_fused_nax( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9)]], // [B, Hv, n_ckpt, Dv, Dk] + device float* chunk_mats [[buffer(10)]], // [B, Hv, n_chunks, 3, 16, 16] + device float* chunk_delta [[buffer(11)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + const int n_chunks = (T + C - 1) / C; + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + auto c_state = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + + auto o_chunk_mats = chunk_mats + n * n_chunks * 3 * 256; + auto o_chunk_delta = chunk_delta + n * n_chunks * C * Dv + dv_idx; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + // Which chunk process_chunk is on, so it knows whether this is a segment + // boundary. + int chunk_idx = 0; + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // Checkpoint the state entering this chunk before anything mutates it, on + // segment boundaries only. + if (save_state_cache && (chunk_idx % Ckpt) == 0) { + S_tile.store(c_state, Dk); + c_state += Dv * Dk; + } + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + float(g_[thread_index_in_simdgroup * Hv + hv_idx]), 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + // KKt_raw is identical for every Dv tile of this chunk; only the first + // tile needs to write it out for the vjp kernel to reuse. + if (save_state_cache && dv_idx == 0) { + KKt_tile.store(o_chunk_mats, 16); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + // Same as KKt_raw: cache the pre-decay inverse before it is overwritten + // below by the gamma-scaled (TUinv) version. + if (save_state_cache && dv_idx == 0) { + Tinv_tile.store(o_chunk_mats + 256, 16); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + QKt_tile.clear(); + tmp_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Hk * Dk); + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); + SCALE_ROW_NAX(Q_tile, gamma); + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); + } + + if (save_state_cache) { + if (dv_idx == 0) { + QKt_tile.store(o_chunk_mats + 512, 16); + } + // Every Dv tile stores its own slice. + delta_tile.store(o_chunk_delta, Dv); + o_chunk_mats += 3 * 256; + o_chunk_delta += C * Dv; + } + + // Output path runs unconditionally: the forward serves both inference and + // the training save, so y and state_out are always needed. + SCALE_TRI_NAX(QKt_tile, gamma) + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + + chunk_idx++; + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp b/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp new file mode 100644 index 00000000..bf46ed67 --- /dev/null +++ b/Source/Cmlx/mlx-generated/gated_delta_update_nax_vjp.cpp @@ -0,0 +1,3816 @@ +namespace mlx::core::metal { + +const char* gated_delta_update_nax_vjp() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/defines.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/defines.h" +// Copyright © 2024 Apple Inc. + + +#define STEEL_CONST static constant constexpr const +#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") +#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/type_traits.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/type_traits.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace metal { + +template +struct is_empty : metal::bool_constant<__is_empty(T)> {}; + +#ifdef __cpp_variable_templates +template +constexpr constant bool is_empty_v = is_empty::value; +#endif + +template +struct make_void { + typedef void type; +}; + +template +using void_t = typename make_void::type; + +template +struct is_static : metal::bool_constant>::value> {}; + +template +struct pointer_element {}; + +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; + +template +using pointer_element_t = typename pointer_element>::type; + +} // namespace metal + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils/integral_constant.h" +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// Integral constant with casting +/////////////////////////////////////////////////////////////////////////////// + +template +struct integral_constant { + static constexpr constant T value = v; + using value_type = T; + using type = integral_constant; + + METAL_FUNC constexpr operator value_type() const thread noexcept { + return value; + } +}; + +template +using bool_constant = integral_constant; +using true_type = bool_constant; +using false_type = bool_constant; + +template +struct is_integral : bool_constant::value> {}; + +template +struct is_integral> + : bool_constant::value> {}; + +template +constexpr constant bool is_integral_v = is_integral::value; + +template +using Int = integral_constant; + +/////////////////////////////////////////////////////////////////////////////// +// Binary Operators on Integral constants +/////////////////////////////////////////////////////////////////////////////// + +#define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ + } + +integral_const_binop(+, operator+); +integral_const_binop(-, operator-); +integral_const_binop(*, operator*); +integral_const_binop(/, operator/); + +integral_const_binop(==, operator==); +integral_const_binop(!=, operator!=); +integral_const_binop(<, operator<); +integral_const_binop(>, operator>); +integral_const_binop(<=, operator<=); +integral_const_binop(>=, operator>=); + +integral_const_binop(&&, operator&&); +integral_const_binop(||, operator||); + +template >> +METAL_FUNC constexpr auto operator||(true_type, T) { + return true_type{}; +} +template >> +METAL_FUNC constexpr auto operator||(T, true_type) { + return true_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(false_type, T) { + return false_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(T, false_type) { + return false_type{}; +} + +// Dispatch utilities +template +void dispatch_bool(bool v, F f) { + if (v) { + f(true_type{}); + } else { + f(false_type{}); + } +} + +template +constexpr void const_for_loop(F f) { + if constexpr (start < stop) { + constexpr auto idx = Int{}; + f(idx); + const_for_loop(f); + } +} + +#undef integral_const_binop + +/////////////////////////////////////////////////////////////////////////////// +// Reduction operators +/////////////////////////////////////////////////////////////////////////////// + +template +METAL_FUNC constexpr T sum(T x) { + return x; +} + +template +METAL_FUNC constexpr auto sum(T x, Us... us) { + return x + sum(us...); +} + +} // namespace steel +} // namespace mlx + +#pragma METAL internals : disable + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/nax.h" +// Copyright © 2025 Apple Inc. + + +#include +#include +#include + + +#include + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// MMA helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// NAX Steel with new tiles +/////////////////////////////////////////////////////////////////////////////// + +struct BaseNAXFrag { + STEEL_CONST short kFragRows = 16; + STEEL_CONST short kFragCols = 16; + + STEEL_CONST short kElemsPerFrag = (kFragRows * kFragCols) / 32; + + STEEL_CONST short kElemRows = 2; + STEEL_CONST short kElemCols = 4; + + STEEL_CONST short kElemRowsJump = 8; + + static_assert( + kElemRows * kElemCols == kElemsPerFrag, + "MMAFrag shape is not consistent with MMAFrag size"); + + template + using dtype_frag_t = typename metal::vec; + + METAL_FUNC static short2 get_coord() { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4; + return short2{fn, fm}; + } + + METAL_FUNC static short2 get_coord(short idx) { + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)) + (idx >> 2) * 8; + const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4 + idx % 4; + return short2{fn, fm}; + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + c + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_rows( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j)]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } + } + + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void load_safe( + thread dtype_frag_t& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((r < lx) && ((c + j) < ly)) { + dst[i * kElemCols + j] = + static_cast(src[r * str_x + (c + j) * str_y]); + } else { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_rows( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + if (r < lx) { + if constexpr (metal::is_same_v>) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + c + j] = static_cast(src[i * kElemCols + j]); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_safe( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = {}, + OffY off_y = {}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; + + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if (r < lx && (c + j) < ly) { + dst[r * str_x + (c + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename T, + typename DstPtrType, + typename StrX, + typename StrY, + typename StartX, + typename StopX, + typename StartY, + typename StopY, + typename OffX = Int<0>, + typename OffY = Int<0>> + METAL_FUNC static constexpr void store_slice( + const thread dtype_frag_t& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + StartX start_x, + StopX stop_x, + StartY start_y, + StopY stop_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + using U = pointer_element_t; + + const short2 sc = get_coord(); + + const_for_loop<0, kElemRows, 1>([&](auto idx_row) { + const auto r = off_x + idx_row * Int{}; + if (r >= stop_x - sc.y || r < start_x - sc.y) { + return; + } + + const_for_loop<0, kElemCols, 1>([&](auto idx_col) { + const auto c = off_y + idx_col; + if (c >= stop_y - sc.x || c < start_y - sc.x) { + return; + } + + const auto src_idx = idx_row * Int{} + idx_col; + dst[(r + sc.y) * str_x + (c + sc.x) * str_y] = + static_cast(src[src_idx]); + }); + }); + } + + template + METAL_FUNC static constexpr void row_reduce( + thread const dtype_frag_t& inp_vals, + thread T* reduced_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + T thr_reduce = Op::apply( + Op::apply(inp_vals[i * kElemCols + 0], inp_vals[i * kElemCols + 1]), + Op::apply(inp_vals[i * kElemCols + 2], inp_vals[i * kElemCols + 3])); + + T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1)); + qgr_reduce = Op::apply(thr_reduce, qgr_reduce); + + T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8)); + sgr_reduce = Op::apply(qgr_reduce, sgr_reduce); + + reduced_vals[i] = Op::apply(reduced_vals[i], sgr_reduce); + } + } + + template + METAL_FUNC static constexpr void row_bin_op( + thread dtype_frag_t& inp_vals, + thread T* row_vals) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + inp_vals[i * kElemCols + j] = + Op::apply(inp_vals[i * kElemCols + j], row_vals[i]); + } + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; + } + } +}; + +template < + typename T, + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> +struct NAXTile { + using NAXFrag_t = NAXFrag_; + using elem_type = T; + + STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; + STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; + STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; + + STEEL_CONST short kTileRows = kTileRows_; + STEEL_CONST short kTileCols = kTileCols_; + + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; + + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; + + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; + + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; + + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; + + METAL_FUNC NAXTile() thread {} + + METAL_FUNC constexpr void clear() thread { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kNumFrags; ++i) { + val_frags[i] = frag_type(0); + } + } + + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; + } + + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j) const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); + } + + template + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_reduce( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + NAXFrag_t::template row_bin_op( + frag_at(i, j), &vptr[i * kFragThrRows]); + } + } + } + + template + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_rows( + frag_at(), + src, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_safe( + frag_at(), + src, + ld, + Int<1>{}, + src_tile_dims.y, + src_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), + dst, + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), + dst, + ld, + Int<1>{}, + dst_tile_dims.y, + dst_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store_slice( + device U* dst, + const int ld, + const short2 start, + const short2 stop) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_slice( + frag_at(), + dst, + ld, + Int<1>{}, + start.y, + stop.y, + start.x, + stop.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } +}; + +template < + class CTile, + class ATile, + class BTile, + bool transpose_a, + bool transpose_b> +METAL_FUNC void tile_matmad_nax( + thread CTile& C, + thread ATile& A, + metal::bool_constant, + thread BTile& B, + metal::bool_constant) { + // Static checks + constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); + + constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); + + constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); + + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; + + if constexpr (TN == 1 && TM % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm += 2) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } + } + } + } +} + +} // namespace steel +} // namespace mlx + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_nax_ops.h" +// Copyright © 2026 Apple Inc. + +#include + +#include +#include + + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/atomic.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/atomic.h" +// Copyright © 2023 Apple Inc. + + +#include +#include + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// Atomic utils +/////////////////////////////////////////////////////////////////////////////// + +#pragma METAL internals : enable +template +constexpr constant bool is_metal_atomic = _disjunction< + is_same, + is_same, + is_same, + is_same>::value; + +#pragma METAL internals : disable + +template +struct mlx_atomic { + atomic val; +}; + +template +struct mlx_atomic>> { + atomic val; +}; + +/////////////////////////////////////////////////////////////////////////////// +// Native metal atomics +/////////////////////////////////////////////////////////////////////////////// + +template , bool> = true> +METAL_FUNC T +mlx_atomic_load_explicit(device mlx_atomic* object, size_t offset) { + return atomic_load_explicit(&(object[offset].val), memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void +mlx_atomic_store_explicit(device mlx_atomic* object, T val, size_t offset) { + atomic_store_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_and_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_and_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_or_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_or_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_min_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_max_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_add_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + atomic_fetch_add_explicit(&(object[offset].val), val, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_mul_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + T expected = mlx_atomic_load_explicit(object, offset); + while (!mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val * expected, offset)) { + // Workaround: Metal's atomic_compare_exchange_weak_explicit does + // not perform bitwise comparison as required by the C++ atomics spec. + // The compiler lowers the success check to `fcmp fast ueq` under + // no-nans-fp-math, which evaluates to false when either operand is NaN - + // even when the bit patterns are identical. With NaN in memory the CAS + // can never succeed, so the loop spins. Bail out instead: memory is + // already NaN and that is the correct reduction result regardless of + // this thread's update. + if constexpr (metal::is_floating_point_v) { + if (isnan(expected)) { + break; + } + } + } +} + +template , bool> = true> +METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit( + device mlx_atomic* object, + thread T* expected, + T val, + size_t offset) { + return atomic_compare_exchange_weak_explicit( + &(object[offset].val), + expected, + val, + memory_order_relaxed, + memory_order_relaxed); +} + +// Specialization for float since it does not atomic_fetch_min_explicit +template <> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + float val, + size_t offset) { + float expected = mlx_atomic_load_explicit(object, offset); + while (val < expected) { + if (mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val, offset)) { + return; + } + } +} + +// Specialization for float since it does not atomic_fetch_max_explicit +template <> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + float val, + size_t offset) { + float expected = mlx_atomic_load_explicit(object, offset); + while (val > expected) { + if (mlx_atomic_compare_exchange_weak_explicit( + object, &expected, val, offset)) { + return; + } + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Custom atomics +/////////////////////////////////////////////////////////////////////////////// + +namespace { + +template +constexpr constant uint packing_size = sizeof(uint) / sizeof(T); + +template +union uint_or_packed { + T val[packing_size]; + uint bits; +}; + +template +struct mlx_atomic_update_helper { + uint operator()(uint_or_packed init, T update, size_t elem_offset) thread { + Op op; + init.val[elem_offset] = op(update, init.val[elem_offset]); + return init.bits; + } +}; + +template +METAL_FUNC void mlx_atomic_update_and_store( + device mlx_atomic* object, + T update, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + + mlx_atomic_update_helper helper; + uint_or_packed expected; + expected.bits = + atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed); + + while (Op::condition(update, expected.val[elem_offset]) && + !mlx_atomic_compare_exchange_weak_explicit( + object, + &(expected.bits), + helper(expected, update, elem_offset), + pack_offset)) { + } +} + +template +struct __None { + static bool condition(T a, T b) { +#pragma unused(a) +#pragma unused(b) + return true; + } + + T operator()(T a, T b) thread { +#pragma unused(b) + return a; + } +}; + +template +struct __Add { + static bool condition(T a, T b) { +#pragma unused(a) +#pragma unused(b) + return true; + } + + T operator()(T a, T b) thread { + return a + b; + } +}; + +template +struct __Mul { + static bool condition(T a, T b) { +#pragma unused(a) + return b != 0; + } + + T operator()(T a, T b) thread { + return a * b; + } +}; + +template +struct __Max { + static bool condition(T a, T b) { + return a > b; + } + + T operator()(T a, T b) thread { + return max(a, b); + } +}; + +template +struct __Min { + static bool condition(T a, T b) { + return a < b; + } + + T operator()(T a, T b) thread { + return min(a, b); + } +}; + +} // namespace + +template , bool> = true> +METAL_FUNC T +mlx_atomic_load_explicit(device mlx_atomic* object, size_t offset) { + size_t pack_offset = offset / sizeof(T); + size_t elem_offset = offset % sizeof(T); + uint_or_packed packed_val; + packed_val.bits = + atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed); + return packed_val.val[elem_offset]; +} + +template , bool> = true> +METAL_FUNC void +mlx_atomic_store_explicit(device mlx_atomic* object, T val, size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_and_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + uint_or_packed identity; + identity.bits = __UINT32_MAX__; + identity.val[elem_offset] = val; + + atomic_fetch_and_explicit( + &(object[pack_offset].val), identity.bits, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_or_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + size_t pack_offset = offset / packing_size; + size_t elem_offset = offset % packing_size; + uint_or_packed identity; + identity.bits = 0; + identity.val[elem_offset] = val; + + atomic_fetch_or_explicit( + &(object[pack_offset].val), identity.bits, memory_order_relaxed); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_min_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_max_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_add_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC void mlx_atomic_fetch_mul_explicit( + device mlx_atomic* object, + T val, + size_t offset) { + mlx_atomic_update_and_store>(object, val, offset); +} + +template , bool> = true> +METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit( + device mlx_atomic* object, + thread uint* expected, + uint val, + size_t offset) { + return atomic_compare_exchange_weak_explicit( + &(object[offset].val), + expected, + val, + memory_order_relaxed, + memory_order_relaxed); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h" +// Copyright © 2026 Apple Inc. + + +#include + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +template +METAL_FUNC void reduce_tile_tg( + thread TileT& acc, + threadgroup float* scratch, + device mlx_atomic* out, + int row_stride, + int col_off, + short valid_rows, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(sg_id * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _s = 1; _s < kNSG; _s++) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(_s * 32 + simd_lane_id) * kE + _i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Simdgroup 0 alone holds the total. + if (sg_id != 0) { + return; + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * row_stride + col_off + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(out, AT_NAX(acc, _i), idx); + } else { + mlx_atomic_store_explicit(out, AT_NAX(acc, _i), idx); + } + } + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + device float* seg_states [[buffer(17)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + const int n_chunks = (T + C - 1) / C; + // Ceiling: the floor form gives 0 when n_chunks < Ckpt, which would leave the + // allocation empty and fault on the first read. + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + const short tail = short(T - (n_chunks - 1) * C); + + // Bases, indexed by chunk rather than walked. With two nested loops the + // incremental decrements are no longer expressible. + auto q_base = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_base = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dq_base = dq + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dk_base = dk + b_idx * T * Hk * Dk + hk_idx * Dk; + + auto v_base = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto dv_base = dv + b_idx * T * Hv * Dv + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto g_base = g + b_idx * T * Hv; + auto beta_base = beta + b_idx * T * Hv; + auto dg_base = dg + b_idx * T * Hv; + auto db_base = db + b_idx * T * Hv; + + auto cm_base = chunk_mats + n * n_chunks * 3 * 256; + auto cd_base = chunk_delta + n * n_chunks * C * Dv + dv_idx; + auto ck_base = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + // Each simdgroup owns a distinct dv slice, so these regions are disjoint and + // the replay needs no cross-simdgroup synchronisation. + auto seg_base = seg_states + (n * Ckpt * Dv + dv_idx) * Dk; + + // Per-chunk pointers, assigned from the bases before each process_chunk call. + auto q_ = q_base; + auto k_ = k_base; + auto dq_ = dq_base; + auto dk_ = dk_base; + auto v_ = v_base; + auto dv_ = dv_base; + auto co_ = co_base; + auto g_ = g_base; + auto beta_ = beta_base; + auto dg_ = dg_base; + auto db_ = db_base; + auto i_chunk_mats = cm_base; + auto i_chunk_delta = cd_base; + auto i_seg = seg_base; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, taken from the replayed segment. + _M16xDk S_tile; + + _M16x32 K_tile, Q_tile; + _M16x16 V_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // From the replayed segment rather than a per-chunk checkpoint. + S_tile.load(i_seg, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + if (fm < valid_rows) { + db_st[fm] = _b0; + } + if (fm1 < valid_rows) { + db_st[fm1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (fm == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (fm < valid_rows) { + dg_st[fm] = _d0; + } + if (fm1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (fm1 < valid_rows) { + dg_st[fm1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk segments in reverse. Each is replayed forward from its checkpoint to + // recover the entry states, then walked backwards. + // Backward over one segment: replay its entry states from the checkpoint, + // then walk it in reverse. HasTail is a compile-time tag so valid_rows is the + // constant C on every chunk of an interior segment, which folds away the + // row guards in the dq/dk stores and the load_rows switch in the replay. + auto process_segment = [&](const int seg, + auto tail_tag) __attribute__((always_inline)) { + constexpr bool HasTail = decltype(tail_tag)::value; + + const int seg_start = seg * Ckpt; + const int rem = n_chunks - seg_start; + const int n_in_seg = (rem < Ckpt) ? rem : Ckpt; + const int n_full = HasTail ? n_in_seg - 1 : n_in_seg; + + // delta is cached per chunk, so the replay is the state recurrence alone: + // S <- exp(gamma_C) S + K_dec^T delta + // No T^-1, no W, no U -- roughly a tenth of a full chunk's work, which is + // what makes trading Ckpt-fold memory for it worthwhile. + // + // Scoped so S is dead before process_chunk runs: it is a _M16xDk, and + // having it live alongside S_tile and dS_tile would put three of them in + // flight. + { + _M16xDk S; + S.load(ck_base + seg * Dv * Dk, Dk); + + auto replay_one = [&](const int j, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + S.store(seg_base + j * Dv * Dk, Dk); + + const int c = seg_start + j; + auto k_c = k_base + c * C * Hk * Dk; + auto g_c = g_base + c * C * Hv; + auto cd_c = cd_base + c * C * Dv; + + float gv = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + g_c[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + auto gs = simd_prefix_inclusive_sum(gv); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gs); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float g_last = gamma[C - 1]; + const float dec[2] = { + metal::fast::exp(g_last - gamma[fm]), + metal::fast::exp(g_last - gamma[fm1])}; + + _M16x16 delta_tile; + delta_tile.load(cd_c, Dv); + + S = S * metal::fast::exp(g_last); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 Kd; + if constexpr (B) { + Kd.load_rows(k_c + kk, Dk * Hk, valid_rows); + } else { + Kd.load(k_c + kk, Dk * Hk); + } + Kd = scale_rows(Kd, dec); + MMA16x32x16(S, kk / 16, delta_tile, true, 0, Kd, false, 0); + } + }; + + for (int j = 0; j < n_full; j++) { + replay_one(j, C, metal::false_type{}); + } + if constexpr (HasTail) { + replay_one(n_in_seg - 1, tail, metal::true_type{}); + } + } + + auto set_chunk_ptrs = [&](const int j) __attribute__((always_inline)) { + const int c = seg_start + j; + q_ = q_base + c * C * Hk * Dk; + k_ = k_base + c * C * Hk * Dk; + dq_ = dq_base + c * C * Hk * Dk; + dk_ = dk_base + c * C * Hk * Dk; + v_ = v_base + c * C * Hv * Dv; + dv_ = dv_base + c * C * Hv * Dv; + co_ = co_base + c * C * Hv * Dv; + g_ = g_base + c * C * Hv; + beta_ = beta_base + c * C * Hv; + dg_ = dg_base + c * C * Hv; + db_ = db_base + c * C * Hv; + i_chunk_mats = cm_base + c * 3 * 256; + i_chunk_delta = cd_base + c * C * Dv; + i_seg = seg_base + j * Dv * Dk; + }; + + if constexpr (HasTail) { + set_chunk_ptrs(n_in_seg - 1); + process_chunk(tail, metal::true_type{}); + } + for (int j = n_full - 1; j >= 0; --j) { + set_chunk_ptrs(j); + process_chunk(C, metal::false_type{}); + } + }; + + // Only the last segment can hold a short chunk + int seg = n_ckpt - 1; + if (tail != C) { + process_segment(seg, metal::true_type{}); + --seg; + } + for (; seg >= 0; --seg) { + process_segment(seg, metal::false_type{}); + } + + dS_tile.store(o_dh, Dk); +} + +template +[[kernel]] void gated_delta_vjp_fused_nax1( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; // KdKb_tmp folded away + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_r0 < valid_rows) { + db_st[_r0] = _b0; + } + if (_r1 < valid_rows) { + db_st[_r1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (_r0 == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (_r0 < valid_rows) { + dg_st[_r0] = _d0; + } + if (_r1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (_r1 < valid_rows) { + dg_st[_r1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +template +METAL_FUNC void reduce_tile_tg0( + thread TileT& acc, + threadgroup float* scratch, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short lo = kNSG / 2; lo > 0; lo >>= 1) { + if (sg_id >= lo && sg_id < 2 * lo) { + const short dst = sg_id - lo; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(dst * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id < lo) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(sg_id * 32 + simd_lane_id) * kE + _i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax0( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 2) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 NI_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(NI_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(NI_tile, _i) = (_c.x == _c.y) ? -1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw; + MUL_NAX(QKt_tile, QKt_tile, tri_decay); + + TUinv_tile = TWinv_tile; + MUL_NAX(TUinv_tile, TUinv_tile, tri_decay); + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + MUL_NAX(dODT, dODT, tri_decay); + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + SCALE2_P(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + _M16x32 Kb_tile; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + + SCALE_BETA_NAX_O(Kb_tile, K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg(dq_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dq_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + const short _fn = _c.x + _f * 16; + const short _fm = _c.y; + if (_fm < valid_rows) { + const int idx = _fm * Hk * Dk + kk + _fn; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } else { + // Dv reduction done in threadgroup memory: single writer. + mlx_atomic_store_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } + } + } + } + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + SCALE_BETA_NAX(dVb_tile, beta_fm); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dVb_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.y < valid_rows) { + dv_[_c.y * Hv * Dv + dv_idx + _c.x] = AT_NAX(dVb_tile, _i); + } + } + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + MUL_NAX(dTU_tile, dTU_tile, tri_decay); // dTinv = dT_W + (dT_U * D). + + ADD_NAX(dTinv_tile, dTW_tile, dTU_tile); + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + SCALE_NAX(dA_tile, -1.0f); + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = dA_tile; + SCALE_TRIEQ_NAX1(G_tile, beta_fm); + _M16x16 GGt_tile; + GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile; + _M16x16 R_tile; + + MUL_NAX(P_tile, QKt_raw, dODT); + MUL_NAX(R_tile, TWinv_tile, dTU_tile); + ADD_NAX(dgam_pair, P_tile, R_tile); + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, dKb_tile, t1, KdKb_tmp; + _M16x32 QgdOS, TbK, KdKC_blk; + + load_seq(Q_tile, q_ + kk, Hk * Dk); // [16 x 32] + load_seq(K_tile, k_ + kk, Dk * Hk); // [16 x 32] + + // gamma: + rowsum(Q * gamma*(dO @ S)) + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(QgdOS, row_exp); + MULA_NAX(dgam_row32, QgdOS, Q_tile); + + // dODT.T @ Q and (G + G.T) @ K + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + // B (T.T @ dW_raw), dW_raw = -gamma * (ddelta @ S) + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + + // dbeta term uses dK_b before the beta scaling, with raw K + MUL_NAX(KdKb_tmp, dKb_tile, K_tile); + ADD_NAX(KdKb, KdKb, KdKb_tmp); + + SCALE_BETA_NAX(dKb_tile, beta_fm); + ADD_NAX(dk_acc, dk_acc, dKb_tile); + + // D_C * (delta @ dS) + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + SCALE2_P(t1, dec_exp); + ADD_NAX(dk_acc, dk_acc, t1); + + // gamma stuff + MUL_NAX(KdKC_blk, K_tile, t1); + SUB_NAX(dgam_row32, dgam_row32, KdKC_blk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC_blk)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC_blk, _i); + } + + // gamma stuff: gamma_C * sum(S * dS) + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16)[_i] * + dS_tile.frag_at(0, kk / 16)[_i]; + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]; + } + + // gamma stuff + SCALE_BETA_NAX(K_tile, beta_fm); + SCALE_NAX(dW_raw, -1.0f); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + MULS_NAX(dgam_row32, TbK, dW_raw); + + // dk reduces over Dv the same way dq does. + reduce_tile_tg(dk_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dk_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * Hk * Dk + kk + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } else { + mlx_atomic_store_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } + } + } + } + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + _M16x16 VdVb; + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + MUL_NAX(VdVb, V_tile, TUddelta); + + _M16x16 AKKt; + TRIL_NAX(AKKt, dA_tile); + MUL_NAX(AKKt, AKKt, KKt_tile); + + ADD_NAX(VdVb, VdVb, AKKt); + ROWSUM1_NAX(dbeta_acc, VdVb); + ROWSUM1_NAX(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _r = 0; _r < 2; _r++) { + float _b = dbeta_acc[_r]; + _b += simd_shuffle_xor(_b, ushort(1)); + _b += simd_shuffle_xor(_b, ushort(8)); + + const short _row = fm + _r * mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_row < valid_rows) { + db_st[_row] = _b; + } + } + } + + // Fold the two fragments: only the row sum of dgam_row is used downstream, + // and row sums are additive across fragments. + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + AT_NAX(dgam_row, _i) = AT_NAX(dgam_row32, _i) + + AT_NAX(dgam_row32, mlx::steel::BaseNAXFrag::kElemsPerFrag + _i); + } + + // reduce and store dgamma: rowsum(row + pair) - colsum(pair) + ADD_NAX(dgam_row, dgam_row, dgam_pair); + _M16x16 dgam_tile, cs_tile; + MM16x16x16(dgam_tile, 0, dgam_row, false, 0, Ones_tile, false, 0); + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + SUB_NAX(dgam_tile, dgam_tile, cs_tile); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dgam_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.x == 0) { + dg_st[_c.y] = AT_NAX(dgam_tile, _i); + } + } + // The gamma_{C-1}-only term lands on the same slot the loop just wrote, + // so the lanes must be ordered before lane 0 accumulates into it. + simdgroup_barrier(mem_flags::mem_threadgroup); + const float dgam_last_red = simd_sum(dgam_last); + if (simd_lane_id == 0 && valid_rows > 0) { + dg_st[valid_rows - 1] += dgam_last_red; + } + } + + // One barrier serves both stages; simdgroup 0 commits. + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + SCALE_ROW_P(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + NSCALE_ROW_P(W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +// Postprocessing dg. Similar idea to what is done in the triton kernel. +// The kernel computes the gradient with respect to the log. +template +[[kernel]] void gated_delta_dgamma_to_dg( + const device InT* g [[buffer(0)]], + device float* dg [[buffer(1)]], // in: dL/dgamma, out: dL/dg + constant int& T [[buffer(2)]], + constant int& Hv [[buffer(3)]], + constant int& n_total [[buffer(4)]], // B * Hv + uint2 pos [[thread_position_in_grid]]) { + const int n = int(pos.x); // b * Hv + hv + const int c = int(pos.y); // chunk index + + if (n >= n_total) { + return; + } + + const int t0 = c * C; + if (t0 >= T) { + return; + } + const int len = min(C, T - t0); + + const int b_idx = n / Hv; + const int hv_idx = n % Hv; + const int base = b_idx * T * Hv + hv_idx; + + float acc = 0.0f; + for (int j = len - 1; j >= 0; --j) { + const int idx = base + (t0 + j) * Hv; + acc += static_cast(dg[idx]); + float gv = metal::max(static_cast(g[idx]), 1e-6f); + dg[idx] = acc / gv; + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/gemv.cpp b/Source/Cmlx/mlx-generated/gemv.cpp index 36ed559b..71a0757c 100644 --- a/Source/Cmlx/mlx-generated/gemv.cpp +++ b/Source/Cmlx/mlx-generated/gemv.cpp @@ -681,10 +681,20 @@ constexpr complex_t operator*(complex_t a, complex_t b) { template constexpr complex_t operator/(complex_t a, complex_t b) { - auto denom = b.real * b.real + b.imag * b.imag; - auto x = a.real * b.real + a.imag * b.imag; - auto y = a.imag * b.real - a.real * b.imag; - return {x / denom, y / denom}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -692,11 +702,7 @@ template < typename U, enable_if_t && is_lane_convertible_v, bool> = true> constexpr complex_t operator/(U a, complex_t b) { - auto scalar = static_cast(a); - auto denom = b.real * b.real + b.imag * b.imag; - auto x = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template @@ -739,6 +745,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: @@ -1133,6 +1140,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -1251,6 +1281,53 @@ inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/utils.h" /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/gemv_masked.cpp b/Source/Cmlx/mlx-generated/gemv_masked.cpp index 8e2faebe..e5269ed6 100644 --- a/Source/Cmlx/mlx-generated/gemv_masked.cpp +++ b/Source/Cmlx/mlx-generated/gemv_masked.cpp @@ -511,7 +511,7 @@ struct GEMVTKernel { const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x); const int out_mask_offset = - !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0]; + !has_output_mask ? 0 : n_block_idx * out_mask_strides[0]; int mat_mask_offset = !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0]; diff --git a/Source/Cmlx/mlx-generated/masked_scatter.cpp b/Source/Cmlx/mlx-generated/masked_scatter.cpp index d660afad..25e4ffb0 100644 --- a/Source/Cmlx/mlx-generated/masked_scatter.cpp +++ b/Source/Cmlx/mlx-generated/masked_scatter.cpp @@ -44,7 +44,7 @@ template if (src_contiguous) { out[idx] = src[batch_idx * src_batch_size + src_index]; } else { - out[idx] = src[elem_to_loc( + out[idx] = src[elem_to_loc( batch_idx * src_batch_size + src_index, src_shapes, src_strides, diff --git a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal index ba0268dc..dd643afe 100644 --- a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal +++ b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal @@ -1,11 +1,17 @@ // Copyright © 2023 Apple Inc. +#include #include #include "utils.h" using namespace metal; +template , bool> = true> +inline bool isnan(T) { + return false; +} + template struct IndexValPair { uint32_t index; @@ -17,8 +23,11 @@ struct ArgMin { static constexpr constant U init = Limits::max; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val > current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val > current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -29,7 +38,7 @@ struct ArgMin { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -43,8 +52,11 @@ struct ArgMax { static constexpr constant U init = Limits::min; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { - if (best.val < current.val || - (best.val == current.val && best.index > current.index)) { + if ((isnan(current.val) && + (!isnan(best.val) || best.index > current.index)) || + (!isnan(best.val) && + (best.val < current.val || + (best.val == current.val && best.index > current.index)))) { return current; } else { return best; @@ -55,7 +67,7 @@ struct ArgMax { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if ((!isnan(best.val) && isnan(vals[i])) || vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/Source/Cmlx/mlx-generated/metal/binary_ops.h b/Source/Cmlx/mlx-generated/metal/binary_ops.h index 37650c7d..968087f8 100644 --- a/Source/Cmlx/mlx-generated/metal/binary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/binary_ops.h @@ -261,7 +261,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/metal/complex.h b/Source/Cmlx/mlx-generated/metal/complex.h index 629b6289..874e837c 100644 --- a/Source/Cmlx/mlx-generated/metal/complex.h +++ b/Source/Cmlx/mlx-generated/metal/complex.h @@ -264,10 +264,20 @@ constexpr complex_t operator*(complex_t a, complex_t b) { template constexpr complex_t operator/(complex_t a, complex_t b) { - auto denom = b.real * b.real + b.imag * b.imag; - auto x = a.real * b.real + a.imag * b.imag; - auto y = a.imag * b.real - a.real * b.imag; - return {x / denom, y / denom}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -275,11 +285,7 @@ template < typename U, enable_if_t && is_lane_convertible_v, bool> = true> constexpr complex_t operator/(U a, complex_t b) { - auto scalar = static_cast(a); - auto denom = b.real * b.real + b.imag * b.imag; - auto x = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template diff --git a/Source/Cmlx/mlx-generated/metal/conv.metal b/Source/Cmlx/mlx-generated/metal/conv.metal index 74027df7..0f49414e 100644 --- a/Source/Cmlx/mlx-generated/metal/conv.metal +++ b/Source/Cmlx/mlx-generated/metal/conv.metal @@ -418,8 +418,8 @@ winograd_conv_2d_weight_transform( // wt_out is stored transposed (A x A x C x O) short ohw_0 = sm * 8 + sn; short ohw_1 = sm * 8 + sn + 1; - device T* wt_out_0 = wt_out + ohw_0 * C * O + ko; - device T* wt_out_1 = wt_out + ohw_1 * C * O + ko; + device T* wt_out_0 = wt_out + size_t(ohw_0) * C * O + ko; + device T* wt_out_1 = wt_out + size_t(ohw_1) * C * O + ko; // Prepare shared memory threadgroup T Ws[BO][R][R][BC]; diff --git a/Source/Cmlx/mlx-generated/metal/cross_entropy.h b/Source/Cmlx/mlx-generated/metal/cross_entropy.h new file mode 100644 index 00000000..bef0bce5 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/cross_entropy.h @@ -0,0 +1,133 @@ +// Copyright © 2026 Apple Inc. + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy( + const device T* in, + const device int32_t* targets, + device float* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + in += gid * size_t(axis_size); + + constexpr int SIMD_SIZE = 32; + + threadgroup AccT local_max[SIMD_SIZE]; + threadgroup AccT local_normalizer[SIMD_SIZE]; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + + AccT prevmax; + AccT maxval = Limits::finite_min; + AccT normalizer = 0; + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + AccT vals[N_READS]; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + vals[i] = AccT(in[offset + i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + vals[i] = + (offset + i < axis_size) ? AccT(in[offset + i]) : Limits::min; + } + } + prevmax = maxval; + for (int i = 0; i < N_READS; i++) { + maxval = (maxval < vals[i]) ? vals[i] : maxval; + } + normalizer *= fast::exp(prevmax - maxval); + for (int i = 0; i < N_READS; i++) { + normalizer += fast::exp(vals[i] - maxval); + } + } + + prevmax = maxval; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + normalizer = simd_sum(normalizer); + + uint n_simdgroups = ceildiv(lsize, uint(SIMD_SIZE)); + prevmax = maxval; + if (simd_lane_id == 0) { + local_max[simd_group_id] = maxval; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + maxval = (simd_lane_id < n_simdgroups) ? local_max[simd_lane_id] + : Limits::finite_min; + maxval = simd_max(maxval); + normalizer *= fast::exp(prevmax - maxval); + if (simd_lane_id == 0) { + local_normalizer[simd_group_id] = normalizer; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + normalizer = + (simd_lane_id < n_simdgroups) ? local_normalizer[simd_lane_id] : AccT(0); + normalizer = simd_sum(normalizer); + + if (lid == 0) { + // Subtract the two logits first; they are the same magnitude, so less is + // lost than in logsumexp - x_t. + AccT gap = maxval - x_t; + out[gid] = isinf(maxval) ? float(gap) : float(log(normalizer) + gap); + } +} + +template < + typename T, + typename AccT = float, + int N_READS = CROSS_ENTROPY_N_READS> +[[kernel]] void cross_entropy_vjp( + const device T* in, + const device int32_t* targets, + const device float* loss, + const device float* cotan, + device T* out, + constant int& axis_size, + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint lsize [[threads_per_threadgroup]]) { + size_t row_offset = gid * size_t(axis_size); + in += row_offset; + out += row_offset; + + int y_n = targets[gid]; + y_n = (y_n < 0) ? y_n + axis_size : y_n; + AccT x_t = AccT(in[y_n]); + AccT g = AccT(cotan[gid]); + AccT l = AccT(loss[gid]); + + // out aliases in when the input is donated, so every thread must read the + // target column before any thread writes. + threadgroup_barrier(mem_flags::mem_device); + + for (int r = 0; r < static_cast(ceildiv(axis_size, N_READS * lsize)); + r++) { + int offset = r * lsize * N_READS + lid * N_READS; + if (offset + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } else { + for (int i = 0; i < N_READS; i++) { + int col = offset + i; + if (col < axis_size) { + AccT p = fast::exp((AccT(in[col]) - x_t) - l); + out[col] = T(g * (p - ((col == y_n) ? AccT(1) : AccT(0)))); + } + } + } + } +} diff --git a/Source/Cmlx/mlx-generated/metal/defines.h b/Source/Cmlx/mlx-generated/metal/defines.h index c369adb7..0e659061 100644 --- a/Source/Cmlx/mlx-generated/metal/defines.h +++ b/Source/Cmlx/mlx-generated/metal/defines.h @@ -14,6 +14,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: diff --git a/Source/Cmlx/mlx-generated/metal/fp8.h b/Source/Cmlx/mlx-generated/metal/fp8.h index 42c5ec12..00efdeb7 100644 --- a/Source/Cmlx/mlx-generated/metal/fp8.h +++ b/Source/Cmlx/mlx-generated/metal/fp8.h @@ -30,11 +30,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized.h b/Source/Cmlx/mlx-generated/metal/fp_quantized.h index 73336195..57d66a7c 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized.h @@ -136,13 +136,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { } template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -154,7 +153,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -164,9 +164,14 @@ struct QuantizedBlockLoader { MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + static_assert( (n_reads * pack_factor) <= group_size, "The number of reads per thread must be less than the group size."); + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); const int src_ld; const int tile_stride; @@ -180,6 +185,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -187,7 +195,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -202,14 +211,19 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales( scales_ + bi * src_ld / group_size + - (bj * pack_factor) / group_size) {} + (bj * pack_factor) / group_size) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -221,21 +235,15 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } return; } - T scale = dequantize_scale(*scales); + float scale = + float(dequantize_scale(*scales)) * inv_scale_enc; for (int i = 0; i < n_reads; i++) { dequantize( src[i * bytes_per_pack], scale, dst + i * pack_factor); @@ -326,6 +334,7 @@ template < typename T, int group_size, int bits, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( @@ -359,6 +368,15 @@ METAL_FUNC void fp_qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // Reuse the last valid row for reads beyond the output tile. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; @@ -368,8 +386,12 @@ METAL_FUNC void fp_qmv_fast_impl( load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device auto* sl = scales + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device auto* sl = scales + src * in_vec_size_g; U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); @@ -387,7 +409,7 @@ METAL_FUNC void fp_qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { if constexpr (has_global_scale) { y[row] = static_cast(result[row] * inv_scale_enc); } else { @@ -788,14 +810,17 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -820,7 +845,7 @@ METAL_FUNC void fp_qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -831,7 +856,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -850,12 +876,15 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_k = group_size % BK != 0; + const int k_blocks = K_eff / BK; + const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0); if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -865,7 +894,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -875,9 +904,18 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, num_els)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -887,7 +925,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K_eff; k += BK) { + for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -898,6 +936,15 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } + if constexpr (partial_k) { + if (num_k > 0) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(num_k, BM)); + loader_w.load_safe(short2(num_k, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + } + } } // Store results to device memory @@ -913,12 +960,14 @@ template < typename T, int group_size, int bits, + bool has_global_scale = false, int BM = 32, int BK = 32, int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Xs, @@ -956,7 +1005,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; auto wl = (const device uint8_t*)w; @@ -971,16 +1021,23 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid, global_scale); mma_t mma_op(simd_gid, simd_lid); + constexpr bool partial_n = group_size % BN != 0; + const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN); + if (num_els < BM) { if ((K % BK) != 0) { const int k_blocks = K / BK; for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -989,9 +1046,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, num_els)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1009,7 +1076,11 @@ METAL_FUNC void fp_qmm_n_impl( for (int k = 0; k < k_blocks; k++) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); - loader_w.load_unsafe(); + if constexpr (partial_n) { + loader_w.load_safe(short2(num_outs, BK)); + } else { + loader_w.load_unsafe(); + } threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); loader_x.next(); @@ -1018,9 +1089,19 @@ METAL_FUNC void fp_qmm_n_impl( const short num_k = K - k_blocks * BK; threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(num_k, BM)); - loader_w.load_safe(short2(BN, num_k)); + loader_w.load_safe(short2(num_outs, num_k)); threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs, Ws); + } else if (num_outs < BN) { + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(num_outs, BK)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } } else { for (int k = 0; k < K; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1036,8 +1117,8 @@ METAL_FUNC void fp_qmm_n_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - if (num_els < BM) { - mma_op.store_result_safe(y, N, short2(BN, num_els)); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); } else { mma_op.store_result(y, N); } @@ -1078,11 +1159,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device uint8_t*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -1125,6 +1207,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -1172,6 +1258,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( @@ -1213,6 +1300,7 @@ template < T, group_size, bits, + partial_rows, has_global_scale, results_per_simdgroup>( w, @@ -1232,6 +1320,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void fp_qmv( @@ -1440,16 +1529,18 @@ template template < typename T, - const int group_size, - const int bits, - const bool aligned_N, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool aligned_N, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -1489,8 +1580,31 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1546,8 +1660,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1576,10 +1690,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1596,7 +1711,7 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( + fp_qmv_fast_impl( w, scales, global_scale, @@ -1635,10 +1750,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1694,10 +1810,11 @@ template uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { int M = x_shape[x_batch_ndims]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1733,12 +1850,14 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> [[kernel]] void fp_gather_qmm_t( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1768,10 +1887,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1788,8 +1908,31 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + T, + has_global_scale, + BM, + BK, + BN>( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1804,7 +1947,7 @@ template < const device uint32_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], const device T* x [[buffer(2)]], - device T* y [[buffer(3)]], + device float* y [[buffer(3)]], const constant int& K [[buffer(4)]], const constant int& N [[buffer(5)]], const constant int& M [[buffer(6)]], @@ -1829,9 +1972,10 @@ template < scales += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - fp_qmm_t_impl( + fp_qmm_t_impl( (const device uint32_t*)wl, scales, + nullptr, x, y, Xs, @@ -1857,6 +2001,7 @@ template < [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1887,10 +2032,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -1907,8 +2053,21 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, + scales, + global_scale, + x, + y, + Xs, + Ws, + K, + N, + M, + tid, + lid, + simd_gid, + simd_lid); } template < @@ -1920,16 +2079,19 @@ template < int BK, int WM, int WN, - bool transpose> + bool transpose, + bool has_global_scale = false> [[kernel]] void fp_gather_qmm_rhs( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1960,7 +2122,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -1973,13 +2136,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1995,112 +2163,61 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h index 35c61c2f..0e4f7ec8 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -60,13 +60,12 @@ struct Dequantize { }; template -inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { - const float s = float(scale); +inline void dequantize(uint8_t w, float scale, threadgroup U* w_local) { if constexpr (bits == 4) { - w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); - w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); + w_local[0] = static_cast(scale * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(scale * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); + w_local[0] = static_cast(scale * Dequantize<8, float>{}(w)); } } @@ -78,7 +77,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool has_global_scale = false> struct QuantizedBlockLoader { MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -106,6 +106,9 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; + // nvfp4 tensor scale, folded into the group scale as fp_dequantize does. + // Kept in float: it is ~1e-5, so in fp16 small scales lose most bits. + float inv_scale_enc = 1.0f; QuantizedBlockLoader( const device uint8_t* src_, @@ -113,7 +116,8 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) thread + ushort simd_lane_id [[thread_index_in_simdgroup]], + const device float* global_scale = nullptr) thread : src_ld(src_ld_), tile_stride( reduction_dim ? BCOLS_PACKED* bytes_per_pack @@ -126,7 +130,11 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id) {} + scales(scales_ + bi * src_ld / group_size + group_id) { + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -135,7 +143,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -149,14 +158,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -165,7 +167,8 @@ struct QuantizedBlockLoader { int k = 0; for (int i = 0; i < n_steps_per_read; i++) { - T scale = dequantize_scale(scales[i]); + float scale = + float(dequantize_scale(scales[i])) * inv_scale_enc; for (int j = 0; j < n_reads_per_scale; j++) { dequantize( src[k * bytes_per_pack], scale, dst + k * pack_factor); @@ -191,6 +194,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -200,6 +204,7 @@ template < METAL_FUNC void fp_qmm_t_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup Wtype* Ws, @@ -229,7 +234,8 @@ METAL_FUNC void fp_qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -245,7 +251,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -335,6 +341,7 @@ template < typename T, const int group_size, const int bits, + const bool has_global_scale = false, const int BM = 64, const int BK = 64, const int BN = 64, @@ -344,6 +351,7 @@ template < METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, threadgroup T* Ws, @@ -373,7 +381,8 @@ METAL_FUNC void fp_qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -391,7 +400,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid, global_scale); constexpr short SM = BM / WM; constexpr short SN = BN / WN; @@ -486,11 +495,12 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +template METAL_FUNC void adjust_matrix_offsets( const device T*& x, const device uint32_t*& w, const device S*& scales, + const device float*& global_scale, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, device T*& y, @@ -533,6 +543,10 @@ METAL_FUNC void adjust_matrix_offsets( w += idx.x; scales += idx.y; } + // One global scale per expert, contiguous over the batch dims of w. + if constexpr (has_global_scale) { + global_scale += w_idx; + } y += tid.z * output_stride; } @@ -547,10 +561,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& K, @@ -589,8 +605,19 @@ template < s_strides, tid); } - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -648,8 +675,8 @@ template < tid); } - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -662,10 +689,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_t_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -694,10 +723,11 @@ template < threadgroup Wtype Ws[BN * BK_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -714,8 +744,19 @@ template < w_strides, s_strides, tid); - fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_t_impl< + T, + group_size, + bits, + aligned_N, + has_global_scale, + BM, + BK, + BN, + WM, + WN, + Wtype>( + w, scales, global_scale, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -727,10 +768,12 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, + const bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_n_nax( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -761,10 +804,11 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - adjust_matrix_offsets( + adjust_matrix_offsets( x, w, scales, + global_scale, lhs_indices, rhs_indices, y, @@ -781,8 +825,8 @@ template < w_strides, s_strides, tid); - fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + fp_qmm_n_impl( + w, scales, nullptr, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -795,16 +839,19 @@ template < int WM, int WN, bool transpose, + bool has_global_scale = false, typename Wtype = bfloat> [[kernel]] void fp_gather_qmm_rhs_nax( const device T* x, const device uint32_t* w, const device uint8_t* scales, - const device uint32_t* indices, + const device float* global_scale, + const device int32_t* offsets, device T* y, const constant int& M, const constant int& N, const constant int& K, + const constant int& num_groups, uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -821,7 +868,8 @@ template < transpose, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + has_global_scale>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -833,13 +881,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -865,10 +918,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -876,149 +929,123 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; - } - } + NAXTile Dtile; + Dtile.clear(); - xn += BK; - loader_w.next(); + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id, + global_scale + group); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); } - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); - - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + xn += BK; + loader_w.next(); + } + + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h b/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h new file mode 100644 index 00000000..154a6e4b --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_nax_ops.h @@ -0,0 +1,588 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include + +#include +#include + +#include "steel/gemm/nax.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +typedef mlx::steel::NAXTile _M16x16; +typedef mlx::steel::NAXTile _M16x32; + +// NAX MACROS +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +// out = a - b +template +METAL_FUNC T operator-(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) - b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator-(const thread T& a) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = -a.frag_at(0, f); + } + return out; +} + +// out = a + b +template +METAL_FUNC T operator+(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator+(const thread T& a, float addend) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) + addend; + } + return out; +} + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +// out = a * b, elementwise. Not a matmul. +template +METAL_FUNC T operator*(const thread T& a, const thread T& b) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T operator*(const thread T& a, float s) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * s; + } + return out; +} + +template +METAL_FUNC T fmadd(const thread T& a, const thread T& b, const thread T& c) { + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = a.frag_at(0, f) * b.frag_at(0, f) + c.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC _M16x16 reduce(const thread T& a) { + _M16x16 out; + out.frag_at(0, 0) = a.frag_at(0, 0); + STEEL_PRAGMA_UNROLL + for (short f = 1; f < T::kNumFrags; f++) { + out.frag_at(0, 0) += a.frag_at(0, f); + } + return out; +} + +template +METAL_FUNC T scale_rows(const thread T& tile, const thread float* s) { + typename T::frag_type sv; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + sv[i] = s[i >> 2]; + } + + T out; + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + out.frag_at(0, f) = tile.frag_at(0, f) * sv; + } + return out; +} + +template +METAL_FUNC void row_sum(thread float* dst, const thread T& tile) { + STEEL_PRAGMA_UNROLL + for (short f = 0; f < T::kNumFrags; f++) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < T::kElemsPerFrag; i++) { + dst[i >> 2] += tile.frag_at(0, f)[i]; + } + } +} + +// Scales row i by DEC2[group(i)] == exp(gamma_{C-1} - gamma_i). +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define TRIL_NAX(TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) = (_c.x >= _c.y) ? 0.f : AT_NAX(TILE1, _i); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + using ct_a_t = decltype(gemm_op.template get_left_input_cooperative_tensor< + AType, + BType, + CType>()); + using ct_b_t = decltype(gemm_op.template get_right_input_cooperative_tensor< + AType, + BType, + CType>()); + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = + gemm_op + .template get_destination_cooperative_tensor(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + typename decltype(C)::elem_type, \ + typename decltype(A)::elem_type, \ + typename decltype(B)::elem_type, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +// Row sum of an elementwise product, reduced over the fragment's column axis. +// DST[0] is the fm row group, DST[1] the fm + kElemRowsJump group. +#define ROWSUM_NAX(DST, TILE0, TILE1) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i) * AT_NAX(TILE1, _i); \ + } \ + } + +// already generic +#define SCALE_BETA_NAX_O(DST, TILE, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(DST, _i) = AT_NAX(TILE, _i) * (BETA2)[_w >> 2]; \ + } \ + } + +#define ROWSUM1_NAX(DST, TILE0) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + (DST)[_w >> 2] += AT_NAX(TILE0, _i); \ + } \ + } + +#define MUL_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULA_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) += AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define MULS_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) -= AT_NAX(TILE1, _i) * AT_NAX(TILE2, _i); \ + } \ + } + +#define NSCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (ROW2)[_w >> 2]; \ + } \ + } + +// As SCALE_ROW_P but also negates, in one pass. +#define NSCALE_ROW_P(TILE0, ROW2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) = -AT_NAX(TILE0, _i) * (ROW2)[_w >> 2]; \ + } \ + } + +#define SCALE2_P(TILE0, DEC2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (DEC2)[_w >> 2]; \ + } \ + } \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h new file mode 100644 index 00000000..841be039 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update.h @@ -0,0 +1,366 @@ +#pragma once + +#include +#include "utils.h" + +constant bool save_state [[function_constant(200)]]; + +#define AT(TILE, IDX) TILE.thread_elements()[IDX] +#define SUB(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) - AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) - AT(TILE2, 1); \ + } +#define ADD(TILE0, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = AT(TILE1, 1) + AT(TILE2, 1); \ + } +#define FMA(TILE0, S, TILE1, TILE2) \ + { \ + AT(TILE0, 0) = S * AT(TILE1, 0) + AT(TILE2, 0); \ + AT(TILE0, 1) = S * AT(TILE1, 1) + AT(TILE2, 1); \ + } + +#define SCALE(TILE0, S) \ + { \ + AT(TILE0, 0) *= S; \ + AT(TILE0, 1) *= S; \ + } +#define SCALE2(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= S0; \ + AT(TILE0, 1) *= S1; \ + } +#define SCALE_TRI(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn > fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 > fm ? 0.f : S1; \ + } +#define SCALE_TRIEQ(TILE0, S0, S1) \ + { \ + AT(TILE0, 0) *= fn >= fm ? 0.f : S0; \ + AT(TILE0, 1) *= fn + 1 >= fm ? 0.f : S1; \ + } + +// lambdas are not supported in metal 14 so porting to macros. + +// non transposed +#define LOAD_M(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fm < valid_rows) ? ((SRC)[fm * (LD) + fn]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fm < valid_rows) ? ((SRC)[fm * (LD) + fn + 1]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fm * (LD) + fn]); \ + AT(M, 1) = static_cast((SRC)[fm * (LD) + fn + 1]); \ + } + +// transposed load: sequence is the column -> mask fn / fn+1 +#define LOAD_MT(M, SRC, LD, B) \ + if constexpr (B) { \ + AT(M, 0) = \ + static_cast((fn < valid_rows) ? ((SRC)[fn * (LD) + fm]) : 0.f); \ + AT(M, 1) = static_cast( \ + (fn + 1 < valid_rows) ? ((SRC)[(fn + 1) * (LD) + fm]) : 0.f); \ + } else { \ + AT(M, 0) = static_cast((SRC)[fn * (LD) + fm]); \ + AT(M, 1) = static_cast((SRC)[(fn + 1) * (LD) + fm]); \ + } + +#define PROCESS_CHUNK_SG(B, S_tile, VALID) \ + { \ + const short valid_rows = (VALID); \ + \ + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) \ + ? metal::fast::log( \ + metal::max( \ + static_cast( \ + g_[thread_index_in_simdgroup * Hv + hv_idx]), \ + 1e-6f)) \ + : 0.0f; \ + \ + float gamma_val = simd_prefix_inclusive_sum(g_val); \ + \ + if (thread_index_in_simdgroup < C) { \ + gamma[thread_index_in_simdgroup] = gamma_val; \ + } \ + simdgroup_barrier(mem_flags::mem_threadgroup); \ + \ + float gamma_fm = metal::fast::exp(gamma[fm]); \ + float gamma_fmdfn = metal::fast::exp(gamma[fm] - gamma[fn]); \ + float gamma_fmdfn1 = metal::fast::exp(gamma[fm] - gamma[fn + 1]); \ + float gamma_Cdfn = metal::fast::exp(gamma[C - 1] - gamma[fn]); \ + float gamma_Cdfn1 = metal::fast::exp(gamma[C - 1] - gamma[fn + 1]); \ + float gamma_C = metal::fast::exp(gamma[C - 1]); \ + \ + float beta_fm = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; \ + \ + KKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + LOAD_MT(KT_tile, k_ + kk, Dk * Hk, B) \ + simdgroup_multiply_accumulate(KKt_tile, K_tile, KT_tile, KKt_tile); \ + } \ + \ + KKtK_tile = KKt_tile; \ + SCALE_TRIEQ(KKtK_tile, beta_fm, beta_fm) \ + \ + simdgroup_float8x8 Tinv, P; \ + AT(P, 0) = AT(KKtK_tile, 0); \ + AT(P, 1) = AT(KKtK_tile, 1); \ + SUB(Tinv, I_tile, KKtK_tile) \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int step = 1; (1 << step) < C; step++) { \ + simdgroup_multiply(P, P, P); \ + simdgroup_multiply_accumulate(Tinv, Tinv, P, Tinv); \ + } \ + \ + WS_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(K_tile, k_ + kk, Dk * Hk, B) \ + SCALE(K_tile, beta_fm) \ + simdgroup_multiply(W_tile, Tinv, K_tile); \ + SCALE(W_tile, gamma_fm) \ + simdgroup_multiply_accumulate(WS_tile, W_tile, S_tile[kk / 8], WS_tile); \ + } \ + \ + SCALE_TRI(Tinv, gamma_fmdfn, gamma_fmdfn1) \ + \ + LOAD_M(V_tile, v_ + dv_idx, Dv * Hv, B) \ + SCALE(V_tile, beta_fm) \ + simdgroup_multiply(U_tile, Tinv, V_tile); \ + SUB(delta_tile, U_tile, WS_tile) \ + \ + tmp_tile = make_filled_simdgroup_matrix(0.f); \ + QKt_tile = make_filled_simdgroup_matrix(0.f); \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_M(Q_tile, q_ + kk, Hk * Dk, B) \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + simdgroup_multiply_accumulate(QKt_tile, Q_tile, K_tile, QKt_tile); \ + SCALE(Q_tile, gamma_fm) \ + simdgroup_multiply_accumulate( \ + tmp_tile, Q_tile, S_tile[kk / 8], tmp_tile); \ + } \ + \ + SCALE_TRI(QKt_tile, gamma_fmdfn, gamma_fmdfn1) \ + \ + simdgroup_multiply_accumulate(out_tile, QKt_tile, delta_tile, tmp_tile); \ + \ + if (fm < valid_rows) { \ + y[fm * Hv * Dv + dv_idx + fn] = static_cast(AT(out_tile, 0)); \ + y[fm * Hv * Dv + dv_idx + fn + 1] = static_cast(AT(out_tile, 1)); \ + } \ + \ + MLX_MTL_PRAGMA_UNROLL \ + for (int kk = 0; kk < Dk; kk += 8) { \ + LOAD_MT(K_tile, k_ + kk, Hk * Dk, B) \ + SCALE2(K_tile, gamma_Cdfn, gamma_Cdfn1) \ + simdgroup_multiply(KD_tile, K_tile, delta_tile); \ + FMA(S_tile[kk / 8], gamma_C, S_tile[kk / 8], KD_tile) \ + } \ + } + +template +[[kernel]] void gated_delta_fused_chunk( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + const short qid = thread_index_in_simdgroup / 4; + const short fm = (qid & 4) + + ((thread_index_in_simdgroup / 2) % 4); // row coordinate of the held tile + const short fn = (qid & 2) * 2 + + (thread_index_in_simdgroup % 2) * 2; // column coordinate of the held tile + + auto dv_idx = thread_position_in_grid.y * 8; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + +#define OUTPUT(T) \ + if (true) { \ + simdgroup_store(T, y); \ + return; \ + } + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + simdgroup_float8x8 S_tile[Dk / 8]; + + // simdgroup matrices + simdgroup_float8x8 V_tile, K_tile, KT_tile, Q_tile; + simdgroup_float8x8 W_tile, U_tile; + simdgroup_float8x8 WS_tile; + simdgroup_float8x8 delta_tile; + simdgroup_float8x8 tmp_tile; + simdgroup_float8x8 QKt_tile; + simdgroup_float8x8 out_tile; + simdgroup_float8x8 KD_tile; + + // tiles for WY form computation + simdgroup_float8x8 KKtK_tile, KKt_tile; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + simdgroup_float8x8 I_tile = make_filled_simdgroup_matrix(0.f); + AT(I_tile, 0) = (fm == fn) ? 1.0f : 0.0f; + AT(I_tile, 1) = (fm == fn + 1) ? 1.0f : 0.0f; + + // load initial state into registers + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_load(S_tile[kk / 8], i_state + kk, Dk, ulong2(0, 0), true); + } + + int t = 0; + MLX_MTL_PRAGMA_UNROLL + for (; t + C <= T; t += C) { + PROCESS_CHUNK_SG(false, S_tile, C); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + PROCESS_CHUNK_SG(true, S_tile, short(T - t)); + } + + MLX_MTL_PRAGMA_UNROLL + for (int kk = 0; kk < Dk; kk += 8) { + simdgroup_store(S_tile[kk / 8], o_state + kk, Dk, ulong2(0, 0), true); + } +} + +/* + auto grid = MTL::Size(32, Dv, B * Hv); + auto threads = MTL::Size(32, 4, 1); + */ +template +[[kernel]] void gated_delta_seq( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], // [B, T, Hv] + const device InT* beta [[buffer(5)]], // [B, T, Hv] + device InT* y [[buffer(6)]], // [B, T, Hv, Dv] + device float* state_out [[buffer(7)]], // [B, Hv, Dv, Dk] + constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9), function_constant(save_state)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + // Checkpoint every Ckpt-th step rather than every step: at T = 4096 with + // B = 4, Hv = 32 a dense cache is 34 GB, which the backward's replay makes + // unnecessary. + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + auto c_state_base = state_cache + n * n_ckpt * Dv * Dk + dv_idx * Dk; + + for (int t = 0; t < T; ++t) { + // Stored before the decay, so slot t/Ckpt holds the state at *entry* to + // step t. The backward replays forward from here, so it must be the entry + // state and not the exit state. + if (save_state && (t % Ckpt) == 0) { + auto c_state = c_state_base + (t / Ckpt) * Dv * Dk; + for (int i = 0; i < n_per_t; ++i) { + c_state[n_per_t * dk_idx + i] = state[i]; + } + } + + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } +} \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal b/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal new file mode 100644 index 00000000..230ef4ca --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update.metal @@ -0,0 +1,50 @@ +#include "gated_delta_update.h" +#include "utils.h" + +using namespace metal; + +#define instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) + +#define instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, c) \ + instantiate_kernel( \ + "gated_delta_fused_chunk_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #c, \ + gated_delta_fused_chunk, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + c) + +// Ckpt is a template parameter of the sequential kernel, so the inference path +// has to name one even though it never stores. These must cover every value +// gated_delta_ckpt() can return, since the host puts it in the kernel name. +#define instantiate_gated_delta_dims(in_type, dk, dv, hk, hv) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gated_delta_update_seq(in_type, dk, dv, hk, hv, 16) \ + instantiate_gated_delta_update_fused_chunk(in_type, dk, dv, hk, hv, 8) + +#define instantiate_gated_delta(in_type) \ + instantiate_gated_delta_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_dims(in_type, 128, 128, 16, 64) + +instantiate_gated_delta(float); +instantiate_gated_delta(bfloat16_t); +instantiate_gated_delta(float16_t); diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h new file mode 100644 index 00000000..d9c85f43 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax.h @@ -0,0 +1,251 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "gated_delta_nax_ops.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +/////////////////////////////////////////////////////////////////////////////// +// Function constants +/////////////////////////////////////////////////////////////////////////////// + +constant bool save_state_cache [[function_constant(200)]]; + +template +[[kernel]] void gated_delta_fused_nax( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + device float* state_cache [[buffer(9)]], // [B, Hv, n_ckpt, Dv, Dk] + device float* chunk_mats [[buffer(10)]], // [B, Hv, n_chunks, 3, 16, 16] + device float* chunk_delta [[buffer(11)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + const int n_chunks = (T + C - 1) / C; + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + auto c_state = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + + auto o_chunk_mats = chunk_mats + n * n_chunks * 3 * 256; + auto o_chunk_delta = chunk_delta + n * n_chunks * C * Dv + dv_idx; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + // Which chunk process_chunk is on, so it knows whether this is a segment + // boundary. + int chunk_idx = 0; + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // Checkpoint the state entering this chunk before anything mutates it, on + // segment boundaries only. + if (save_state_cache && (chunk_idx % Ckpt) == 0) { + S_tile.store(c_state, Dk); + c_state += Dv * Dk; + } + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + float(g_[thread_index_in_simdgroup * Hv + hv_idx]), 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + // KKt_raw is identical for every Dv tile of this chunk; only the first + // tile needs to write it out for the vjp kernel to reuse. + if (save_state_cache && dv_idx == 0) { + KKt_tile.store(o_chunk_mats, 16); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + // Same as KKt_raw: cache the pre-decay inverse before it is overwritten + // below by the gamma-scaled (TUinv) version. + if (save_state_cache && dv_idx == 0) { + Tinv_tile.store(o_chunk_mats + 256, 16); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + QKt_tile.clear(); + tmp_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Hk * Dk); + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); + SCALE_ROW_NAX(Q_tile, gamma); + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); + } + + if (save_state_cache) { + if (dv_idx == 0) { + QKt_tile.store(o_chunk_mats + 512, 16); + } + // Every Dv tile stores its own slice. + delta_tile.store(o_chunk_delta, Dv); + o_chunk_mats += 3 * 256; + o_chunk_delta += C * Dv; + } + + // Output path runs unconditionally: the forward serves both inference and + // the training save, so y and state_out are always needed. + SCALE_TRI_NAX(QKt_tile, gamma) + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + + chunk_idx++; + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h new file mode 100644 index 00000000..0e6ea093 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_nax_vjp.h @@ -0,0 +1,1737 @@ +// Copyright © 2026 Apple Inc. +#pragma once + +#include "gated_delta_nax_ops.h" + +#include +#include "atomic.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +template +METAL_FUNC void reduce_tile_tg( + thread TileT& acc, + threadgroup float* scratch, + device mlx_atomic* out, + int row_stride, + int col_off, + short valid_rows, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(sg_id * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _s = 1; _s < kNSG; _s++) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(_s * 32 + simd_lane_id) * kE + _i]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Simdgroup 0 alone holds the total. + if (sg_id != 0) { + return; + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * row_stride + col_off + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(out, AT_NAX(acc, _i), idx); + } else { + mlx_atomic_store_explicit(out, AT_NAX(acc, _i), idx); + } + } + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + device float* seg_states [[buffer(17)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + const int n_chunks = (T + C - 1) / C; + // Ceiling: the floor form gives 0 when n_chunks < Ckpt, which would leave the + // allocation empty and fault on the first read. + const int n_ckpt = (n_chunks + Ckpt - 1) / Ckpt; + const short tail = short(T - (n_chunks - 1) * C); + + // Bases, indexed by chunk rather than walked. With two nested loops the + // incremental decrements are no longer expressible. + auto q_base = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_base = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dq_base = dq + b_idx * T * Hk * Dk + hk_idx * Dk; + auto dk_base = dk + b_idx * T * Hk * Dk + hk_idx * Dk; + + auto v_base = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto dv_base = dv + b_idx * T * Hv * Dv + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto g_base = g + b_idx * T * Hv; + auto beta_base = beta + b_idx * T * Hv; + auto dg_base = dg + b_idx * T * Hv; + auto db_base = db + b_idx * T * Hv; + + auto cm_base = chunk_mats + n * n_chunks * 3 * 256; + auto cd_base = chunk_delta + n * n_chunks * C * Dv + dv_idx; + auto ck_base = state_cache + (n * n_ckpt * Dv + dv_idx) * Dk; + // Each simdgroup owns a distinct dv slice, so these regions are disjoint and + // the replay needs no cross-simdgroup synchronisation. + auto seg_base = seg_states + (n * Ckpt * Dv + dv_idx) * Dk; + + // Per-chunk pointers, assigned from the bases before each process_chunk call. + auto q_ = q_base; + auto k_ = k_base; + auto dq_ = dq_base; + auto dk_ = dk_base; + auto v_ = v_base; + auto dv_ = dv_base; + auto co_ = co_base; + auto g_ = g_base; + auto beta_ = beta_base; + auto dg_ = dg_base; + auto db_ = db_base; + auto i_chunk_mats = cm_base; + auto i_chunk_delta = cd_base; + auto i_seg = seg_base; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, taken from the replayed segment. + _M16xDk S_tile; + + _M16x32 K_tile, Q_tile; + _M16x16 V_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // From the replayed segment rather than a per-chunk checkpoint. + S_tile.load(i_seg, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + if (fm < valid_rows) { + db_st[fm] = _b0; + } + if (fm1 < valid_rows) { + db_st[fm1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (fm == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (fm < valid_rows) { + dg_st[fm] = _d0; + } + if (fm1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (fm1 < valid_rows) { + dg_st[fm1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk segments in reverse. Each is replayed forward from its checkpoint to + // recover the entry states, then walked backwards. + // Backward over one segment: replay its entry states from the checkpoint, + // then walk it in reverse. HasTail is a compile-time tag so valid_rows is the + // constant C on every chunk of an interior segment, which folds away the + // row guards in the dq/dk stores and the load_rows switch in the replay. + auto process_segment = [&](const int seg, + auto tail_tag) __attribute__((always_inline)) { + constexpr bool HasTail = decltype(tail_tag)::value; + + const int seg_start = seg * Ckpt; + const int rem = n_chunks - seg_start; + const int n_in_seg = (rem < Ckpt) ? rem : Ckpt; + const int n_full = HasTail ? n_in_seg - 1 : n_in_seg; + + // delta is cached per chunk, so the replay is the state recurrence alone: + // S <- exp(gamma_C) S + K_dec^T delta + // No T^-1, no W, no U -- roughly a tenth of a full chunk's work, which is + // what makes trading Ckpt-fold memory for it worthwhile. + // + // Scoped so S is dead before process_chunk runs: it is a _M16xDk, and + // having it live alongside S_tile and dS_tile would put three of them in + // flight. + { + _M16xDk S; + S.load(ck_base + seg * Dv * Dk, Dk); + + auto replay_one = [&](const int j, + const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + S.store(seg_base + j * Dv * Dk, Dk); + + const int c = seg_start + j; + auto k_c = k_base + c * C * Hk * Dk; + auto g_c = g_base + c * C * Hv; + auto cd_c = cd_base + c * C * Dv; + + float gv = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + g_c[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + auto gs = simd_prefix_inclusive_sum(gv); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gs); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float g_last = gamma[C - 1]; + const float dec[2] = { + metal::fast::exp(g_last - gamma[fm]), + metal::fast::exp(g_last - gamma[fm1])}; + + _M16x16 delta_tile; + delta_tile.load(cd_c, Dv); + + S = S * metal::fast::exp(g_last); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 Kd; + if constexpr (B) { + Kd.load_rows(k_c + kk, Dk * Hk, valid_rows); + } else { + Kd.load(k_c + kk, Dk * Hk); + } + Kd = scale_rows(Kd, dec); + MMA16x32x16(S, kk / 16, delta_tile, true, 0, Kd, false, 0); + } + }; + + for (int j = 0; j < n_full; j++) { + replay_one(j, C, metal::false_type{}); + } + if constexpr (HasTail) { + replay_one(n_in_seg - 1, tail, metal::true_type{}); + } + } + + auto set_chunk_ptrs = [&](const int j) __attribute__((always_inline)) { + const int c = seg_start + j; + q_ = q_base + c * C * Hk * Dk; + k_ = k_base + c * C * Hk * Dk; + dq_ = dq_base + c * C * Hk * Dk; + dk_ = dk_base + c * C * Hk * Dk; + v_ = v_base + c * C * Hv * Dv; + dv_ = dv_base + c * C * Hv * Dv; + co_ = co_base + c * C * Hv * Dv; + g_ = g_base + c * C * Hv; + beta_ = beta_base + c * C * Hv; + dg_ = dg_base + c * C * Hv; + db_ = db_base + c * C * Hv; + i_chunk_mats = cm_base + c * 3 * 256; + i_chunk_delta = cd_base + c * C * Dv; + i_seg = seg_base + j * Dv * Dk; + }; + + if constexpr (HasTail) { + set_chunk_ptrs(n_in_seg - 1); + process_chunk(tail, metal::true_type{}); + } + for (int j = n_full - 1; j >= 0; --j) { + set_chunk_ptrs(j); + process_chunk(C, metal::false_type{}); + } + }; + + // Only the last segment can hold a short chunk + int seg = n_ckpt - 1; + if (tail != C) { + process_segment(seg, metal::true_type{}); + --seg; + } + for (; seg >= 0; --seg) { + process_segment(seg, metal::false_type{}); + } + + dS_tile.store(o_dh, Dk); +} + +template +[[kernel]] void gated_delta_vjp_fused_nax1( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 1) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + _M16x16 tril_mask; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tril_mask)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); + AT_NAX(tril_mask, _i) = (_c.x >= _c.y) ? 0.0f : 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw * tri_decay; + TUinv_tile = TWinv_tile * tri_decay; + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + dODT = dODT * tri_decay; + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + K_tile = scale_rows(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + _M16x32 Kb_tile = scale_rows(K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dq_acc = scale_rows(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg( + dq_acc, + red_scratch, + dq_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + dVb_tile = scale_rows(dVb_tile, beta_fm); + + dVb_tile.store_rows(dv_ + dv_idx, Hv * Dv, valid_rows); + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + V_tile = scale_rows(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + dTU_tile = dTU_tile * tri_decay; + + dTinv_tile = dTW_tile + dTU_tile; + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + dA_tile = dA_tile * -1.0f; + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = scale_rows(dA_tile, beta_fm) * tril_mask; + _M16x16 GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile = QKt_raw * dODT; + _M16x16 R_tile = TWinv_tile * dTU_tile; + dgam_pair = P_tile + R_tile; + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, t1; + + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + { + _M16x32 QgdOS; + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + dgam_row32 = dgam_row32 + scale_rows(QgdOS, row_exp) * Q_tile; + } + + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + dW_raw = scale_rows(-dW_raw, row_exp); + + { + _M16x32 dKb_tile; + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + KdKb = KdKb + dKb_tile * K_tile; // KdKb_tmp folded away + dk_acc = dk_acc + scale_rows(dKb_tile, beta_fm); + } + + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + t1 = scale_rows(t1, dec_exp); + dk_acc = dk_acc + t1; + + { + const _M16x32 KdKC = K_tile * t1; + dgam_row32 = dgam_row32 - KdKC; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC, _i); + } + } + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * + (S_tile.frag_at(0, kk / 16)[_i] * dS_tile.frag_at(0, kk / 16)[_i] + + S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]); + } + + { + _M16x32 TbK; + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + dgam_row32 = fmadd(TbK, dW_raw, dgam_row32); + } + + reduce_tile_tg( + dk_acc, + red_scratch, + dk_, + Hk * Dk, + kk, + valid_rows, + sg_id, + simd_lane_id); + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + + _M16x16 AKKt = dA_tile * KKt_tile; + AKKt = AKKt * tril_mask; + _M16x16 VdVb = fmadd(V_tile, TUddelta, AKKt); + row_sum(dbeta_acc, VdVb); + row_sum(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + float _b0 = dbeta_acc[0]; + float _b1 = dbeta_acc[1]; + _b0 += simd_shuffle_xor(_b0, ushort(1)); + _b0 += simd_shuffle_xor(_b0, ushort(8)); + _b1 += simd_shuffle_xor(_b1, ushort(1)); + _b1 += simd_shuffle_xor(_b1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_r0 < valid_rows) { + db_st[_r0] = _b0; + } + if (_r1 < valid_rows) { + db_st[_r1] = _b1; + } + } + + dgam_row = reduce(dgam_row32); + + float dgam_acc[2] = {0.0f, 0.0f}; + dgam_row = dgam_row + dgam_pair; + row_sum(dgam_acc, dgam_row); + + _M16x16 cs_tile; + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + + const float dgam_last_red = simd_sum(dgam_last); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + float _d0 = dgam_acc[0]; + float _d1 = dgam_acc[1]; + _d0 += simd_shuffle_xor(_d0, ushort(1)); + _d0 += simd_shuffle_xor(_d0, ushort(8)); + _d1 += simd_shuffle_xor(_d1, ushort(1)); + _d1 += simd_shuffle_xor(_d1, ushort(8)); + + const short _r0 = fm; + const short _r1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + + _d0 -= AT_NAX(cs_tile, 0); + _d1 -= AT_NAX(cs_tile, mlx::steel::BaseNAXFrag::kElemCols); + + if (_r0 == valid_rows - 1) { + _d0 += dgam_last_red; + } + if (_r0 < valid_rows) { + dg_st[_r0] = _d0; + } + if (_r1 == valid_rows - 1) { + _d1 += dgam_last_red; + } + if (_r1 < valid_rows) { + dg_st[_r1] = _d1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + Q_tile = scale_rows(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + K_tile = scale_rows(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + W_raw = scale_rows(-W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +template +METAL_FUNC void reduce_tile_tg0( + thread TileT& acc, + threadgroup float* scratch, + const short sg_id, + const ushort simd_lane_id) { + constexpr short kE = TileT::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short lo = kNSG / 2; lo > 0; lo >>= 1) { + if (sg_id >= lo && sg_id < 2 * lo) { + const short dst = sg_id - lo; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + scratch[(dst * 32 + simd_lane_id) * kE + _i] = AT_NAX(acc, _i); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id < lo) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < kE; _i++) { + AT_NAX(acc, _i) += scratch[(sg_id * 32 + simd_lane_id) * kE + _i]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +[[kernel]] void gated_delta_vjp_fused_nax0( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* beta [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B, Hv, n_chunks, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device float* dg [[buffer(12)]], + device float* db [[buffer(13)]], + device float* dh [[buffer(14)]], + // Cached, chunk-local forward intermediates produced by the forward-save + // pass (see gated_delta_fused_nax): avoids redoing the K@K^T contraction, + // the 15-step Neumann inversion, and the W/U/S contraction for delta. + const device float* chunk_mats [[buffer(15)]], + const device float* chunk_delta [[buffer(16)]], // [B, Hv, n_chunks, C, Dv] + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + using _M16xDk = mlx::steel::NAXTile; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + const int n_chunks = (T + C - 1) / C; + const int t_last = (n_chunks - 1) * C; + + // Pointers positioned at the final chunk + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dq_ = dq + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + auto dk_ = dk + b_idx * T * Hk * Dk + hk_idx * Dk + t_last * Hk * Dk; + + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto dv_ = dv + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + auto co_ = cot_o + b_idx * T * Hv * Dv + hv_idx * Dv + t_last * Hv * Dv; + + auto g_ = g + b_idx * T * Hv + t_last * Hv; + auto beta_ = beta + b_idx * T * Hv + t_last * Hv; + auto dg_ = dg + b_idx * T * Hv + t_last * Hv; + auto db_ = db + b_idx * T * Hv + t_last * Hv; + + auto c_state = state_cache + (n * n_chunks * Dv + dv_idx) * Dk + + (n_chunks - 1) * Dv * Dk; + + // Cached forward intermediates, walked in reverse alongside c_state. + auto i_chunk_mats = + chunk_mats + n * n_chunks * 3 * 256 + (n_chunks - 1) * 3 * 256; + auto i_chunk_delta = + chunk_delta + n * n_chunks * C * Dv + dv_idx + (n_chunks - 1) * C * Dv; + + auto i_cot_h = cot_h + (n * Dv + dv_idx) * Dk; + auto o_dh = dh + (n * Dv + dv_idx) * Dk; + + // One threadgroup per (b, hv) now covers every Dv slice, so the Dv reduction + // never leaves threadgroup memory. + constexpr int kNSG = Dv / 16; + // Under GQA several threadgroups still share an hk, so dq/dk keep an atomic + // for that axis -- but with kNSG writers folded away first. + constexpr bool kGQA = (Hv != Hk); + + threadgroup float gamma_all[C * kNSG]; + threadgroup float* gamma = gamma_all + sg_id * C; + + // (kNSG/2) destinations x 32 lanes x kElemsPerTile(_M16x32) floats. + threadgroup float red_scratch[(kNSG / 2) * 32 * 16]; + threadgroup float db_stage[kNSG * C]; + threadgroup float dg_stage[kNSG * C]; + + float beta_fm[2]; + + // Carried state gradient, dL/dS. Same [dv, dk] orientation as S_tile. + _M16xDk dS_tile; + dS_tile.load(i_cot_h, Dk); + + // Forward state at chunk entry, reloaded from the checkpoint each chunk + _M16xDk S_tile; + + // Recomputed forward tiles + _M16x32 K_tile, Q_tile; + _M16xDk W_tile; + _M16x16 K16_tile, Q16_tile; + _M16x16 V_tile; + _M16x16 U_tile; + _M16x16 WS_tile; + _M16x16 delta_tile; + _M16x16 QKt_tile, QKt_raw; + _M16x16 KKt_tile; + _M16x16 TWinv_tile, TUinv_tile; + _M16x16 TMP_tile; + + // Backward tiles + _M16x16 dout_tile; + _M16x16 ddelta_tile; + _M16x16 dTU_tile, dTW_tile; + _M16x16 dTinv_tile; + _M16x16 dA_tile, G_tile; + _M16x16 tri_decay; + + _M16x16 I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(I_tile, _i) = (_c.x == _c.y) ? 1.0f : 0.0f; + } + + _M16x16 NI_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(NI_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + AT_NAX(NI_tile, _i) = (_c.x == _c.y) ? -1.0f : 0.0f; + } + + _M16x16 Ones_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(Ones_tile)::kElemsPerFrag; _i++) { + AT_NAX(Ones_tile, _i) = 1.0f; + } + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + // reload the checkpoint + S_tile.load(c_state, Dk); + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max(g_[thread_index_in_simdgroup * Hv + hv_idx], 1e-6)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + const float gamma_last = gamma[C - 1]; + const float gamma_last_exp = metal::fast::exp(gamma_last); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + // Per-lane decay factors. gamma[j] == gamma[valid_rows-1] for + // j >= valid_rows (the prefix sum of zeros), so the tail chunk needs no + // special case here. + const float row_exp[2] = { + metal::fast::exp(gamma[fm]), metal::fast::exp(gamma[fm1])}; + const float dec_exp[2] = { + metal::fast::exp(gamma_last - gamma[fm]), + metal::fast::exp(gamma_last - gamma[fm1])}; + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(tri_decay)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + AT_NAX(tri_decay, _i) = + (_c.x > _c.y) ? 0.f : metal::fast::exp(gamma[_c.y] - gamma[_c.x]); + } + + KKt_tile.load(i_chunk_mats, 16); + TWinv_tile.load(i_chunk_mats + 256, 16); + QKt_raw.load(i_chunk_mats + 512, 16); + + QKt_tile = QKt_raw; + MUL_NAX(QKt_tile, QKt_tile, tri_decay); + + TUinv_tile = TWinv_tile; + MUL_NAX(TUinv_tile, TUinv_tile, tri_decay); + + delta_tile.load(i_chunk_delta, Dv); + + // dgamma accumulators: + // dgam_row : per-row + // dgam_pair : row-minus-col + // dgam_last : the gamma_{C-1} + _M16x16 dgam_row, dgam_pair; + dgam_row.clear(); + dgam_pair.clear(); + float dgam_last = 0.0f; + float dbeta_acc[2] = {0.0f, 0.0f}; + + load_seq(dout_tile, co_ + dv_idx, Hv * Dv); + + // dODT = (dO @ delta.T * D) + _M16x16 dODT; + MM16x16x16(dODT, 0, dout_tile, false, 0, delta_tile, true, 0); + MUL_NAX(dODT, dODT, tri_decay); + + // cannot fuse this because ddelta is used by the next loop + MM16x16x16(ddelta_tile, 0, QKt_tile, true, 0, dout_tile, false, 0); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + SCALE2_P(K_tile, dec_exp); + MMA16x16x32(ddelta_tile, 0, K_tile, false, 0, dS_tile, true, kk / 16); + } + + // ddelta = M^T @ dO + K_dec @ dS^T + // dq = gamma * (dO @ S) + dODT @ K + dTW_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dW_raw; + _M16x32 dq_acc; + _M16x32 Kb_tile; + load_seq(K_tile, k_ + kk, Dk * Hk); + + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + + SCALE_BETA_NAX_O(Kb_tile, K_tile, beta_fm); + MMA16x16x32(dTW_tile, 0, dW_raw, false, 0, Kb_tile, true, 0); + + MM16x32x16(dq_acc, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(dq_acc, row_exp); + MMA16x32x16(dq_acc, 0, dODT, false, 0, K_tile, false, 0); + + // dq reduces over Dv, which is now entirely inside this threadgroup. + reduce_tile_tg(dq_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dq_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + const short _fn = _c.x + _f * 16; + const short _fm = _c.y; + if (_fm < valid_rows) { + const int idx = _fm * Hk * Dk + kk + _fn; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } else { + // Dv reduction done in threadgroup memory: single writer. + mlx_atomic_store_explicit(dq_, AT_NAX(dq_acc, _i), idx); + } + } + } + } + } + + // dv = B(Tu.T @ ddelta). Indexed by dv_idx, so no reduction. + _M16x16 dVb_tile; + MM16x16x16(dVb_tile, 0, TUinv_tile, true, 0, ddelta_tile, false, 0); + _M16x16 TUddelta = dVb_tile; // pre-beta, reused by dbeta + SCALE_BETA_NAX(dVb_tile, beta_fm); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dVb_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.y < valid_rows) { + dv_[_c.y * Hv * Dv + dv_idx + _c.x] = AT_NAX(dVb_tile, _i); + } + } + + // dT_U = ddelta @ (beta * V).T + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(dTU_tile, 0, ddelta_tile, false, 0, V_tile, true, 0); + MUL_NAX(dTU_tile, dTU_tile, tri_decay); // dTinv = dT_W + (dT_U * D). + + ADD_NAX(dTinv_tile, dTW_tile, dTU_tile); + + // dA = -T.T @ dTinv @ T.T + MM16x16x16(TMP_tile, 0, TWinv_tile, true, 0, dTinv_tile, false, 0); + MM16x16x16(dA_tile, 0, TMP_tile, false, 0, TWinv_tile, true, 0); + SCALE_NAX(dA_tile, -1.0f); + + // G = tril_(dA) * beta, and GGt = G + G.T + G_tile = dA_tile; + SCALE_TRIEQ_NAX1(G_tile, beta_fm); + _M16x16 GGt_tile; + GGt_tile = G_tile; + MMA16x16x16(GGt_tile, 0, I_tile, true, 0, G_tile, true, 0); + + // dgamma pair sites + _M16x16 P_tile; + _M16x16 R_tile; + + MUL_NAX(P_tile, QKt_raw, dODT); + MUL_NAX(R_tile, TWinv_tile, dTU_tile); + ADD_NAX(dgam_pair, P_tile, R_tile); + + _M16x32 KdKb; // rowsum(K * dK_b) source, needed for dbeta + KdKb.clear(); + _M16x32 dgam_row32; + dgam_row32.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 dk_acc, dW_raw, dKb_tile, t1, KdKb_tmp; + _M16x32 QgdOS, TbK, KdKC_blk; + + load_seq(Q_tile, q_ + kk, Hk * Dk); // [16 x 32] + load_seq(K_tile, k_ + kk, Dk * Hk); // [16 x 32] + + // gamma: + rowsum(Q * gamma*(dO @ S)) + MM16x32x16(QgdOS, 0, dout_tile, false, 0, S_tile, false, kk / 16); + SCALE_ROW_P(QgdOS, row_exp); + MULA_NAX(dgam_row32, QgdOS, Q_tile); + + // dODT.T @ Q and (G + G.T) @ K + MM16x32x16(dk_acc, 0, dODT, true, 0, Q_tile, false, 0); + MMA16x32x16(dk_acc, 0, GGt_tile, false, 0, K_tile, false, 0); + + // B (T.T @ dW_raw), dW_raw = -gamma * (ddelta @ S) + MM16x32x16(dW_raw, 0, ddelta_tile, false, 0, S_tile, false, kk / 16); + NSCALE_ROW_P(dW_raw, row_exp); + MM16x32x16(dKb_tile, 0, TWinv_tile, true, 0, dW_raw, false, 0); + + // dbeta term uses dK_b before the beta scaling, with raw K + MUL_NAX(KdKb_tmp, dKb_tile, K_tile); + ADD_NAX(KdKb, KdKb, KdKb_tmp); + + SCALE_BETA_NAX(dKb_tile, beta_fm); + ADD_NAX(dk_acc, dk_acc, dKb_tile); + + // D_C * (delta @ dS) + MM16x32x16(t1, 0, delta_tile, false, 0, dS_tile, false, kk / 16); + SCALE2_P(t1, dec_exp); + ADD_NAX(dk_acc, dk_acc, t1); + + // gamma stuff + MUL_NAX(KdKC_blk, K_tile, t1); + SUB_NAX(dgam_row32, dgam_row32, KdKC_blk); + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(KdKC_blk)::kElemsPerTile; _i++) { + dgam_last += AT_NAX(KdKC_blk, _i); + } + + // gamma stuff: gamma_C * sum(S * dS) + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16)[_i] * + dS_tile.frag_at(0, kk / 16)[_i]; + dgam_last += gamma_last_exp * S_tile.frag_at(0, kk / 16 + 1)[_i] * + dS_tile.frag_at(0, kk / 16 + 1)[_i]; + } + + // gamma stuff + SCALE_BETA_NAX(K_tile, beta_fm); + SCALE_NAX(dW_raw, -1.0f); + MM16x32x16(TbK, 0, TWinv_tile, false, 0, K_tile, false, 0); + MULS_NAX(dgam_row32, TbK, dW_raw); + + // dk reduces over Dv the same way dq does. + reduce_tile_tg(dk_acc, red_scratch, sg_id, simd_lane_id); + + if (sg_id == 0) { + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dk_acc)::kElemsPerTile; _i++) { + const short _f = _i / mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_w); // {fn, fm} + if (_c.y < valid_rows) { + const int idx = _c.y * Hk * Dk + kk + _c.x + _f * 16; + if (kGQA) { + mlx_atomic_fetch_add_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } else { + mlx_atomic_store_explicit(dk_, AT_NAX(dk_acc, _i), idx); + } + } + } + } + } + + // dbeta = rowsum(V * dV_b) + rowsum(K * dK_b) + rowsum(tril_(dA) * KKt) + _M16x16 VdVb; + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + MUL_NAX(VdVb, V_tile, TUddelta); + + _M16x16 AKKt; + TRIL_NAX(AKKt, dA_tile); + MUL_NAX(AKKt, AKKt, KKt_tile); + + ADD_NAX(VdVb, VdVb, AKKt); + ROWSUM1_NAX(dbeta_acc, VdVb); + ROWSUM1_NAX(dbeta_acc, KdKb); + + { + threadgroup float* db_st = db_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _r = 0; _r < 2; _r++) { + float _b = dbeta_acc[_r]; + _b += simd_shuffle_xor(_b, ushort(1)); + _b += simd_shuffle_xor(_b, ushort(8)); + + const short _row = fm + _r * mlx::steel::BaseNAXFrag::kElemRowsJump; + if (_row < valid_rows) { + db_st[_row] = _b; + } + } + } + + // Fold the two fragments: only the row sum of dgam_row is used downstream, + // and row sums are additive across fragments. + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { + AT_NAX(dgam_row, _i) = AT_NAX(dgam_row32, _i) + + AT_NAX(dgam_row32, mlx::steel::BaseNAXFrag::kElemsPerFrag + _i); + } + + // reduce and store dgamma: rowsum(row + pair) - colsum(pair) + ADD_NAX(dgam_row, dgam_row, dgam_pair); + _M16x16 dgam_tile, cs_tile; + MM16x16x16(dgam_tile, 0, dgam_row, false, 0, Ones_tile, false, 0); + MM16x16x16(cs_tile, 0, dgam_pair, true, 0, Ones_tile, false, 0); + SUB_NAX(dgam_tile, dgam_tile, cs_tile); + + { + threadgroup float* dg_st = dg_stage + sg_id * C; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(dgam_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + if (_c.x == 0) { + dg_st[_c.y] = AT_NAX(dgam_tile, _i); + } + } + // The gamma_{C-1}-only term lands on the same slot the loop just wrote, + // so the lanes must be ordered before lane 0 accumulates into it. + simdgroup_barrier(mem_flags::mem_threadgroup); + const float dgam_last_red = simd_sum(dgam_last); + if (simd_lane_id == 0 && valid_rows > 0) { + dg_st[valid_rows - 1] += dgam_last_red; + } + } + + // One barrier serves both stages; simdgroup 0 commits. + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg_id == 0 && thread_index_in_simdgroup < (uint)valid_rows) { + float b_sum = 0.0f; + float g_sum = 0.0f; + STEEL_PRAGMA_UNROLL + for (short _s = 0; _s < kNSG; _s++) { + b_sum += db_stage[_s * C + thread_index_in_simdgroup]; + g_sum += dg_stage[_s * C + thread_index_in_simdgroup]; + } + const int idx = thread_index_in_simdgroup * Hv + hv_idx; + db_[idx] = b_sum; + dg_[idx] = g_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // dS update + // dS = gamma_C * dS + dO.T @ (gamma * Q) - ddelta.T @ W + SCALE_NAX(dS_tile, gamma_last_exp); + for (int kk = 0; kk < Dk; kk += 32) { + _M16x32 term, W_raw; + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Dk * Hk); + + // dO.T @ (gamma * Q) -> [dv, dk] + SCALE_ROW_P(Q_tile, row_exp); + MM16x32x16(term, 0, dout_tile, true, 0, Q_tile, false, 0); + + // - ddelta.T @ W, with W = gamma * (TWinv @ beta*K) + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_raw, 0, TWinv_tile, false, 0, K_tile, false, 0); + NSCALE_ROW_P(W_raw, row_exp); + MMA16x32x16(term, 0, ddelta_tile, true, 0, W_raw, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(term)::kElemsPerFrag; _i++) { + dS_tile.frag_at(0, kk / 16)[_i] += AT_NAX(term, _i); + dS_tile.frag_at(0, kk / 16 + 1)[_i] += + AT_NAX(term, decltype(term)::kElemsPerFrag + _i); + } + } + }; + + // Walk chunks in reverse. The tail chunk comes first. + int c = n_chunks - 1; + const short tail = short(T - c * C); + if (tail != C) { + process_chunk(tail, metal::true_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + --c; + } + for (; c >= 0; --c) { + process_chunk(C, metal::false_type{}); + q_ -= C * Hk * Dk; + k_ -= C * Hk * Dk; + v_ -= C * Hv * Dv; + co_ -= C * Hv * Dv; + g_ -= C * Hv; + beta_ -= C * Hv; + dq_ -= C * Hk * Dk; + dk_ -= C * Hk * Dk; + dv_ -= C * Hv * Dv; + dg_ -= C * Hv; + db_ -= C * Hv; + c_state -= Dv * Dk; + i_chunk_mats -= 3 * 256; + i_chunk_delta -= C * Dv; + } + + dS_tile.store(o_dh, Dk); +} + +// Postprocessing dg. Similar idea to what is done in the triton kernel. +// The kernel computes the gradient with respect to the log. +template +[[kernel]] void gated_delta_dgamma_to_dg( + const device InT* g [[buffer(0)]], + device float* dg [[buffer(1)]], // in: dL/dgamma, out: dL/dg + constant int& T [[buffer(2)]], + constant int& Hv [[buffer(3)]], + constant int& n_total [[buffer(4)]], // B * Hv + uint2 pos [[thread_position_in_grid]]) { + const int n = int(pos.x); // b * Hv + hv + const int c = int(pos.y); // chunk index + + if (n >= n_total) { + return; + } + + const int t0 = c * C; + if (t0 >= T) { + return; + } + const int len = min(C, T - t0); + + const int b_idx = n / Hv; + const int hv_idx = n % Hv; + const int base = b_idx * T * Hv + hv_idx; + + float acc = 0.0f; + for (int j = len - 1; j >= 0; --j) { + const int idx = base + (t0 + j) * Hv; + acc += static_cast(dg[idx]); + float gv = metal::max(static_cast(g[idx]), 1e-6f); + dg[idx] = acc / gv; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h new file mode 100644 index 00000000..86082c27 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.h @@ -0,0 +1,184 @@ +// Copyright © 2024 Apple Inc. +#pragma once + +#include +#include + +#include "atomic.h" +#include "utils.h" + +using namespace metal; + +template +[[kernel]] void gated_delta_vjp_seq( + const device InT* q [[buffer(0)]], // [B, T, Hk, Dk] + const device InT* k [[buffer(1)]], // [B, T, Hk, Dk] + const device InT* v [[buffer(2)]], // [B, T, Hv, Dv] + const device InT* g [[buffer(3)]], // [B, T, Hv] + const device InT* b [[buffer(4)]], // [B, T, Hv] + const device InT* cot_o [[buffer(5)]], // [B, T, Hv, Dv] + const device float* cot_h [[buffer(6)]], // [B, Hv, Dv, Dk] + const device float* state_cache [[buffer(7)]], // [B*Hv, n_ckpt, Dv, Dk] + constant int& T [[buffer(8)]], + device mlx_atomic* dq [[buffer(9)]], + device mlx_atomic* dk [[buffer(10)]], + device float* dv [[buffer(11)]], + device mlx_atomic* dg [[buffer(12)]], + device mlx_atomic* db [[buffer(13)]], + device float* dh [[buffer(14)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + const int qk_stride = Hk * Dk; + const int v_stride = Hv * Dv; + const int g_stride = Hv; + + auto q_base = q + b_idx * T * qk_stride + hk_idx * Dk; + auto k_base = k + b_idx * T * qk_stride + hk_idx * Dk; + auto dq_base = dq + b_idx * T * qk_stride + hk_idx * Dk; + auto dk_base = dk + b_idx * T * qk_stride + hk_idx * Dk; + + auto v_base = v + b_idx * T * v_stride + hv_idx * Dv; + auto dv_base = dv + b_idx * T * v_stride + hv_idx * Dv; + auto co_base = cot_o + b_idx * T * v_stride + hv_idx * Dv; + + auto g_base = g + b_idx * T * g_stride + hv_idx; + auto b_base = b + b_idx * T * g_stride + hv_idx; + auto dg_base = dg + b_idx * T * g_stride + hv_idx; + auto db_base = db + b_idx * T * g_stride + hv_idx; + + const int n_ckpt = (T + Ckpt - 1) / Ckpt; + + float s_hat[n_per_t]; + auto base_state = cot_h + (n * Dv + dv_idx) * Dk; + for (int i = 0; i < n_per_t; i++) { + s_hat[i] = base_state[n_per_t * dk_idx + i]; + } + + // Only this thread's n_per_t slice of the state is ever needed, so a whole + // segment of entry states costs Ckpt * n_per_t registers and the replay needs + // no device buffer. + float seg[Ckpt][n_per_t]; + float s_prev[n_per_t]; + float s_dec[n_per_t]; + + for (int seg_idx = n_ckpt - 1; seg_idx >= 0; --seg_idx) { + const int t0 = seg_idx * Ckpt; + const int seg_len = metal::min(Ckpt, T - t0); + + // Replay forward from the checkpoint, recording the entry state of each + // step in the segment. + auto c_state = + state_cache + n * n_ckpt * Dv * Dk + seg_idx * Dv * Dk + dv_idx * Dk; + + float s[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + s[i] = c_state[n_per_t * dk_idx + i]; + } + + for (int j = 0; j < seg_len; ++j) { + const int t = t0 + j; + + for (int i = 0; i < n_per_t; ++i) { + seg[j][i] = s[i]; + } + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto k_t = k_base + t * qk_stride; + + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + const int s_idx = n_per_t * dk_idx + i; + s[i] *= gamma; + kv_mem += s[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + + float delta = + beta * (static_cast(v_base[t * v_stride + dv_idx]) - kv_mem); + + for (int i = 0; i < n_per_t; ++i) { + s[i] += delta * static_cast(k_t[n_per_t * dk_idx + i]); + } + } + + // Backward over the same steps, newest first. + for (int j = seg_len - 1; j >= 0; --j) { + const int t = t0 + j; + + float gamma = static_cast(g_base[t * g_stride]); + float beta = static_cast(b_base[t * g_stride]); + + auto q_t = q_base + t * qk_stride; + auto k_t = k_base + t * qk_stride; + auto dq_t = dq_base + t * qk_stride; + auto dk_t = dk_base + t * qk_stride; + + float kv_mem = 0.0f; + float co = static_cast(co_base[t * v_stride + dv_idx]); + float w = 0.0f; + for (int i = 0; i < n_per_t; i++) { + const int s_idx = n_per_t * dk_idx + i; + s_prev[i] = seg[j][i]; + s_dec[i] = s_prev[i] * gamma; + kv_mem += s_dec[i] * static_cast(k_t[s_idx]); + + s_hat[i] += co * static_cast(q_t[s_idx]); + + w += s_hat[i] * static_cast(k_t[s_idx]); + } + kv_mem = simd_sum(kv_mem); + w = simd_sum(w); + + if (thread_index_in_simdgroup == 0) { + dv_base[t * v_stride + dv_idx] = beta * w; + } + + float u = static_cast(v_base[t * v_stride + dv_idx]) - kv_mem; + float delta = beta * u; + + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(db_base + t * g_stride, w * u, 0); + } + + float dgamma = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + + float s_t = s_dec[i] + delta * static_cast(k_t[s_idx]); + mlx_atomic_fetch_add_explicit(dq_t, co * s_t, s_idx); + + float contrib = beta * (u * s_hat[i] - w * s_dec[i]); + mlx_atomic_fetch_add_explicit(dk_t, contrib, s_idx); + + s_hat[i] -= beta * w * static_cast(k_t[s_idx]); + + dgamma += s_hat[i] * s_prev[i]; + } + dgamma = simd_sum(dgamma); + if (thread_index_in_simdgroup == 0) { + mlx_atomic_fetch_add_explicit(dg_base + t * g_stride, dgamma, 0); + } + + for (int i = 0; i < n_per_t; ++i) { + s_hat[i] *= gamma; + } + } + } + + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + dh[(n * Dv + dv_idx) * Dk + s_idx] = s_hat[i]; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal new file mode 100644 index 00000000..172cef32 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gated_delta_update_vjp.metal @@ -0,0 +1,35 @@ +#include "gated_delta_update_vjp.h" +#include "gated_delta_update.h" +#include "utils.h" + +using namespace metal; + +#define instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, ckpt) \ + instantiate_kernel( \ + "seq_gated_delta_vjp_" #in_type "_" #dk "_" #dv "_" #hk "_" #hv \ + "_" #ckpt, \ + gated_delta_vjp_seq, \ + in_type, \ + dk, \ + dv, \ + hk, \ + hv, \ + ckpt) + +#define instantiate_gated_delta_vjp_dims(in_type, dk, dv, hk, hv) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 1) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 4) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 8) \ + instantiate_gdn_vjp_seq(in_type, dk, dv, hk, hv, 16) + +#define instantiate_gated_delta_vjp(in_type) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 24, 24) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 32, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 32) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 48) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 16) \ + instantiate_gated_delta_vjp_dims(in_type, 128, 128, 16, 64) + +instantiate_gated_delta_vjp(float); +instantiate_gated_delta_vjp(bfloat16_t); +instantiate_gated_delta_vjp(float16_t); \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal b/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal new file mode 100644 index 00000000..40b933d5 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/gather_mm_offsets.metal @@ -0,0 +1,23 @@ +// Copyright © 2026 Apple Inc. + +#include + +using namespace metal; + +[[kernel]] void gather_mm_offsets( + const device uint32_t* indices [[buffer(0)]], + device int32_t* offsets [[buffer(1)]], + const constant int& M [[buffer(2)]], + uint g [[thread_position_in_grid]]) { + int lo = 0; + int hi = M; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (indices[mid] < g) { + lo = mid + 1; + } else { + hi = mid; + } + } + offsets[g] = lo; +} diff --git a/Source/Cmlx/mlx-generated/metal/gemv_masked.h b/Source/Cmlx/mlx-generated/metal/gemv_masked.h index 30371eb0..930ba3b9 100644 --- a/Source/Cmlx/mlx-generated/metal/gemv_masked.h +++ b/Source/Cmlx/mlx-generated/metal/gemv_masked.h @@ -452,7 +452,7 @@ struct GEMVTKernel { const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x); const int out_mask_offset = - !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0]; + !has_output_mask ? 0 : n_block_idx * out_mask_strides[0]; int mat_mask_offset = !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0]; diff --git a/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h b/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h index 2ba54740..fe2e2b2f 100644 --- a/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h +++ b/Source/Cmlx/mlx-generated/metal/indexing/masked_scatter.h @@ -32,7 +32,7 @@ template if (src_contiguous) { out[idx] = src[batch_idx * src_batch_size + src_index]; } else { - out[idx] = src[elem_to_loc( + out[idx] = src[elem_to_loc( batch_idx * src_batch_size + src_index, src_shapes, src_strides, diff --git a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h index 3a97de9b..e361e741 100644 --- a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h +++ b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h @@ -38,8 +38,8 @@ METAL_FUNC void scatter_impl( for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) { LocT out_idx = out_offset; for (int i = 0; i < NIDX; ++i) { - auto idx_loc = indices.row_contiguous[i] - ? ind_idx + LocT idx_loc = indices.row_contiguous[i] + ? LocT(ind_idx) : elem_to_loc( ind_idx, &indices.shapes[indices.ndim * i], @@ -50,7 +50,7 @@ METAL_FUNC void scatter_impl( out_idx += static_cast(idx_val) * static_cast(out_strides[ax]); } - auto upd_idx = ind_idx * static_cast(upd_size) + gid.x; + LocT upd_idx = ind_idx * static_cast(upd_size) + gid.x; if constexpr (!UPD_ROW_CONTIG) { upd_idx = elem_to_loc(upd_idx, upd_shape, upd_strides, upd_ndim); } diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 42855e14..48fe55cd 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -28,15 +28,30 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -46,8 +61,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -61,7 +76,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -71,8 +86,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -86,7 +101,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -107,15 +122,30 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -125,8 +155,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -141,7 +171,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -151,8 +181,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -166,7 +196,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -196,13 +226,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -298,13 +342,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -395,11 +453,25 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + result[8 * i] += x * (select(U(0), scale, bool(wb & 0x01)) + bias); + result[8 * i + 1] += x * (select(U(0), scale, bool(wb & 0x02)) + bias); + result[8 * i + 2] += x * (select(U(0), scale, bool(wb & 0x04)) + bias); + result[8 * i + 3] += x * (select(U(0), scale, bool(wb & 0x08)) + bias); + result[8 * i + 4] += x * (select(U(0), scale, bool(wb & 0x10)) + bias); + result[8 * i + 5] += x * (select(U(0), scale, bool(wb & 0x20)) + bias); + result[8 * i + 6] += x * (select(U(0), scale, bool(wb & 0x40)) + bias); + result[8 * i + 7] += x * (select(U(0), scale, bool(wb & 0x80)) + bias); + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -485,14 +557,36 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { template inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -585,9 +679,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -596,6 +690,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -651,14 +751,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -754,7 +847,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -766,7 +859,8 @@ METAL_FUNC void qmv_fast_impl( uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { - constexpr int packs_per_thread = bits == 2 ? 1 : 2; + constexpr int packs_per_thread = + bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor(); @@ -788,19 +882,37 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; y += tid.x * out_vec_size + out_row; - for (int k = 0; k < in_vec_size; k += block_size) { + const int aligned_end = (in_vec_size / block_size) * block_size; + + for (int k = 0; k < aligned_end; k += block_size) { U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -813,9 +925,37 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } + // Tail block: handles K that is a multiple of 512 (dispatch gate) but not + // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for + // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; + // their result[row] stays at 0 and contributes 0 to the simd_sum below. + if (aligned_end < in_vec_size) { + bool in_bounds = + (aligned_end + static_cast(simd_lid) * values_per_thread) < + in_vec_size; + if (in_bounds) { + U sum = load_vector(x, x_thread); + + for (int row = 0; row < results_per_simdgroup; row++) { + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; + + U s = sl[0]; + U b = bl[0]; + result[row] += + qdot(wl, x_thread, s, b, sum); + } + } + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1194,6 +1334,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1202,7 +1343,7 @@ METAL_FUNC void qmm_t_impl( const device T* scales, const device T* biases, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -1227,7 +1368,7 @@ METAL_FUNC void qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1603,6 +1744,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( @@ -1643,7 +1785,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1661,6 +1803,7 @@ template < int group_size, const int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv( @@ -1897,6 +2040,7 @@ template < const int bits, const bool aligned_N, const bool batched, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1946,7 +2090,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -1977,7 +2121,7 @@ template < const device T* scales [[buffer(1)]], const device T* biases [[buffer(2)]], const device T* x [[buffer(3)]], - device T* y [[buffer(4)]], + device float* y [[buffer(4)]], const constant int& K [[buffer(5)]], const constant int& N [[buffer(6)]], const constant int& M [[buffer(7)]], @@ -2005,7 +2149,7 @@ template < biases += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2276,6 +2420,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -2335,7 +2480,7 @@ template < s_strides, b_strides, tid); - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2437,11 +2582,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2485,13 +2631,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2508,113 +2659,61 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } @@ -2662,14 +2761,23 @@ template w_min = simd_min(w_min); w_max = simd_max(w_max); - float scale = max((w_max - w_min) / n_bins, eps); - bool side = abs(w_min) > abs(w_max); - scale = side ? scale : -scale; - float edge = side ? w_min : w_max; - float q0 = round(edge / scale); - bool at_zero = q0 == 0.0f; - scale = at_zero ? scale : edge / q0; - float bias = at_zero ? 0 : edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = max(w_max - w_min, eps); + bias = w_min; + } else { + scale = max((w_max - w_min) / n_bins, eps); + bool side = abs(w_min) > abs(w_max); + scale = side ? scale : -scale; + float edge = side ? w_min : w_max; + float q0 = round(edge / scale); + bool at_zero = q0 == 0.0f; + scale = at_zero ? scale : edge / q0; + bias = at_zero ? 0 : edge; + } // Write out the scales and biases size_t gindex = in_index / group_size; @@ -2773,7 +2881,9 @@ template #pragma clang loop unroll(full) for (int i = 0; i < pack_factor; i++) { uint8_t d; - if (bits == 2) { + if (bits == 1) { + d = (val >> i) & 0x01; + } else if (bits == 2) { d = (val >> (bits * i)) & 0x03; } else if (bits == 4) { d = (val >> (bits * i)) & 0x0f; diff --git a/Source/Cmlx/mlx-generated/metal/quantized_nax.h b/Source/Cmlx/mlx-generated/metal/quantized_nax.h index ed32eb59..95f26868 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/quantized_nax.h @@ -31,15 +31,30 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -49,8 +64,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -64,7 +79,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -74,8 +89,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -89,7 +104,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -110,15 +125,30 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -128,8 +158,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -144,7 +174,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -154,8 +184,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -169,7 +199,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -199,13 +229,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -301,13 +345,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -398,11 +456,33 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + U s[8] = { + scale, + scale / 2.0f, + scale / 4.0f, + scale / 8.0f, + scale / 16.0f, + scale / 32.0f, + scale / 64.0f, + scale / 128.0f}; + for (int i = 0; i < (values_per_thread / 8); i++) { + result[8 * i] += x * (s[0] * (w[i] & 0x01) + bias); + result[8 * i + 1] += x * (s[1] * (w[i] & 0x02) + bias); + result[8 * i + 2] += x * (s[2] * (w[i] & 0x04) + bias); + result[8 * i + 3] += x * (s[3] * (w[i] & 0x08) + bias); + result[8 * i + 4] += x * (s[4] * (w[i] & 0x10) + bias); + result[8 * i + 5] += x * (s[5] * (w[i] & 0x20) + bias); + result[8 * i + 6] += x * (s[6] * (w[i] & 0x40) + bias); + result[8 * i + 7] += x * (s[7] * (w[i] & 0x80) + bias); + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -487,14 +567,36 @@ template inline void dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -587,9 +689,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -653,14 +755,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -722,9 +817,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -793,14 +888,7 @@ struct QuantizedBlockLoader< return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1215,7 +1303,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1334,7 +1423,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_gather_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1477,11 +1567,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1510,13 +1601,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1543,12 +1639,11 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1556,145 +1651,119 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); + NAXTile Dtile; + Dtile.clear(); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); + } - volatile int compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } - } - xn += BK; - loader_w.next(); - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); + (void)compiler_barrier; + } + } - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + xn += BK; + loader_w.next(); + } - volatile int compiler_barrier; + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); + threadgroup_barrier(mem_flags::mem_threadgroup); - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + } - threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); - } + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h index b1546adb..fce36a4c 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h @@ -117,7 +117,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); @@ -172,6 +172,7 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], + const constant int64_t& column_offset [[buffer(11)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -193,7 +194,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT out_idx = gid.y + gsize.y * IdxT(gid.z); @@ -243,7 +244,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += out_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -311,7 +312,8 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], - const constant size_t& out_size [[buffer(11)]], + const constant int64_t& column_offset [[buffer(11)]], + const constant size_t& out_size [[buffer(12)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -336,7 +338,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT full_idx = gid.y + gsize.y * IdxT(gid.z); @@ -383,7 +385,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += full_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h index 604efa78..9a1e363e 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_init.h @@ -3,6 +3,7 @@ template [[kernel]] void init_reduce( device T* out [[buffer(0)]], - uint tid [[thread_position_in_grid]]) { - out[tid] = Op::init; + uint2 index [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + out[index.x + grid_dim.x * int64_t(index.y)] = Op::init; } diff --git a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal index 8fdd0876..a27ad455 100644 --- a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal @@ -2,10 +2,28 @@ // clang-format off #include "utils.h" +#include "sdpa_blocked.h" #include "sdpa_vector.h" using namespace metal; +instantiate_kernel( + "sdpa_blocked_scale_copy_bfloat16", + sdpa_blocked_scale_copy, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_scale_copy_float16", + sdpa_blocked_scale_copy, + float16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_bfloat16", + sdpa_blocked_causal_mask, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_float16", + sdpa_blocked_causal_mask, + float16_t) + // SDPA vector instantiations #define instantiate_sdpa_vector_aggregation(type, value_dim) \ instantiate_kernel( \ @@ -28,32 +46,50 @@ using namespace metal; qk_dim, \ value_dim) -#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ - instantiate_kernel( \ - "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ - sdpa_vector_2pass_1_gqa, \ - type, \ - qk_dim, \ - value_dim, \ - 8, \ +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, g, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #g "_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + g, \ hpt) #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ + instantiate_sdpa_vector(type, 96, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ - instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ - instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ + instantiate_sdpa_vector(type, 512, 512) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 8, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 12, 4) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 16, 2) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ instantiate_sdpa_vector_aggregation(type, 192) \ - instantiate_sdpa_vector_aggregation(type, 256) + instantiate_sdpa_vector_aggregation(type, 256) \ + instantiate_sdpa_vector_aggregation(type, 512) instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) instantiate_sdpa_vector_heads(float16_t) - // clang-format on + +#define instantiate_sdpa_vector_tk1(type, dim) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_" #type "_" #dim "_" #dim "_tk1", \ + sdpa_vector_2pass_1, \ + type, \ + dim, \ + dim, \ + 1) + +instantiate_sdpa_vector_tk1(bfloat16_t, 256) +instantiate_sdpa_vector_tk1(float, 192) +instantiate_sdpa_vector_tk1(float, 256) +// clang-format on diff --git a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal new file mode 100644 index 00000000..b96809e5 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention_vjp.metal @@ -0,0 +1,154 @@ +// Copyright © 2026 Apple Inc. + +#include "utils.h" + +using namespace metal; + +template +[[kernel]] void sdpa_vjp_odo( + const device InT* o [[buffer(0)]], // [B, H, qL, Dv] + const device InT* cot_o [[buffer(1)]], // [B, H, qL, Dv] + device float* odo [[buffer(2)]], // [B, H, qL] + constant int& qL [[buffer(3)]], + uint3 tpg [[thread_position_in_grid]], + uint simd_lane_id [[thread_index_in_simdgroup]]) { + const uint bh = tpg.z; + const int i = int(tpg.y); + + if (i >= qL) { + return; + } + + const size_t row = (size_t)bh * qL + i; + auto o_i = o + row * Dv; + auto co_i = cot_o + row * Dv; + + float acc = 0.0f; + for (int d = int(simd_lane_id); d < Dv; d += 32) { + acc += float(o_i[d]) * float(co_i[d]); + } + acc = simd_sum(acc); + + if (simd_lane_id == 0) { + odo[row] = acc; + } +} + +struct SDPAVJPTileParams { + int bq; + int bk; + int qL; + int i0; + int j0; + float scale; + int diag_off; + int causal; +}; + +template +[[kernel]] void sdpa_vjp_ds( + const device T* S [[buffer(0)]], // [BH, bq, bk] + const device T* dP [[buffer(1)]], // [BH, bq, bk] + const device float* lse [[buffer(2)]], // [B, H, qL] + const device float* odo [[buffer(3)]], // [B, H, qL] + device T* dS [[buffer(4)]], // [BH, bq, bk] + device T* P [[buffer(5)]], // [BH, bk, bq] + const constant SDPAVJPTileParams& p [[buffer(6)]], + uint3 gid [[thread_position_in_grid]]) { + int col = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (row >= p.bq || col >= p.bk) { + return; + } + + size_t sbase = (size_t(bh) * size_t(p.bq) + size_t(row)) * size_t(p.bk); + size_t qi = size_t(bh) * size_t(p.qL) + size_t(p.i0 + row); + + float l = lse[qi]; + float dlt = odo[qi]; + // Fully masked rows carry lse == -inf; exp(-inf - -inf) would be NaN. + bool dead = (l < 0.0f) && metal::isinf(l); + int lim = p.i0 + row + p.diag_off; + + float pv; + if (dead || (p.causal != 0 && (p.j0 + col) > lim)) { + pv = 0.0f; + } else { + pv = metal::fast::exp(static_cast(S[sbase + col]) * p.scale - l); + } + float dsv = pv * (static_cast(dP[sbase + col]) - dlt) * p.scale; + + dS[sbase + col] = static_cast(dsv); + P[sbase + col] = static_cast(pv); +} + +template +[[kernel]] void sdpa_vjp_reduce( + const device T* src [[buffer(0)]], + device float* acc [[buffer(1)]], + device T* out [[buffer(2)]], + const constant int& rows [[buffer(3)]], + const constant int& dim [[buffer(4)]], + const constant int& group [[buffer(5)]], + const constant int& acc_rows [[buffer(6)]], + const constant int& row_off [[buffer(7)]], + uint3 gid [[thread_position_in_grid]]) { + int c = int(gid.x); + int row = int(gid.y); + int bh = int(gid.z); + if (c >= dim || row >= rows) { + return; + } + + size_t sbase = size_t(bh) * size_t(group) * size_t(rows) * size_t(dim); + float sum = 0.0f; + for (int g = 0; g < group; ++g) { + size_t si = sbase + + (size_t(g) * size_t(rows) + size_t(row)) * size_t(dim) + size_t(c); + sum += static_cast(src[si]); + } + + size_t ai = + (size_t(bh) * size_t(acc_rows) + size_t(row_off + row)) * size_t(dim) + + size_t(c); + float total = Accum ? acc[ai] + sum : sum; + acc[ai] = total; + out[ai] = static_cast(total); +} + +// Instantiations +#define instantiate_odo(in_type, dv) \ + instantiate_kernel( \ + "sdpa_vjp_odo_" #in_type "_" #dv, sdpa_vjp_odo, in_type, dv) + +#define instantiate_odo_shapes(in_type) \ + instantiate_odo(in_type, 64); \ + instantiate_odo(in_type, 72); \ + instantiate_odo(in_type, 80); \ + instantiate_odo(in_type, 96); \ + instantiate_odo(in_type, 128); \ + instantiate_odo(in_type, 192); \ + instantiate_odo(in_type, 256); \ + instantiate_odo(in_type, 512); \ + +instantiate_odo_shapes(bfloat16_t); +instantiate_odo_shapes(float16_t); +instantiate_odo_shapes(float); + +#define instantiate_sdpa_vjp_ds(tname, type) \ + instantiate_kernel("sdpa_vjp_ds_" #tname, sdpa_vjp_ds, type) + +instantiate_sdpa_vjp_ds(float, float); +instantiate_sdpa_vjp_ds(float16_t, float16_t); +instantiate_sdpa_vjp_ds(bfloat16_t, bfloat16_t); + +#define instantiate_sdpa_vjp_reduce(tname, type) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_add_" #tname, sdpa_vjp_reduce, type, true) \ + instantiate_kernel( \ + "sdpa_vjp_reduce_set_" #tname, sdpa_vjp_reduce, type, false) + +instantiate_sdpa_vjp_reduce(float, float); +instantiate_sdpa_vjp_reduce(float16_t, float16_t); +instantiate_sdpa_vjp_reduce(bfloat16_t, bfloat16_t); diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h b/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h new file mode 100644 index 00000000..444d5e9e --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/sdpa_blocked.h @@ -0,0 +1,45 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "defines.h" + +using namespace metal; + +template +[[kernel]] void sdpa_blocked_scale_copy( + const device T* q [[buffer(0)]], + device T* q_block [[buffer(1)]], + constant float& scale [[buffer(2)]], + constant int64_t* q_strides [[buffer(3)]], + constant int& heads [[buffer(4)]], + constant int& rows [[buffer(5)]], + constant int& head_dim [[buffer(6)]], + constant int& q_offset [[buffer(7)]], + uint3 elem [[thread_position_in_grid]]) { + int d = elem.x; + int row = elem.y; + int head = elem.z % heads; + int batch = elem.z / heads; + int64_t q_index = int64_t(batch) * q_strides[0] + + int64_t(head) * q_strides[1] + int64_t(q_offset + row) * q_strides[2] + d; + size_t q_block_index = (size_t(elem.z) * rows + row) * head_dim + d; + q_block[q_block_index] = q[q_index] * static_cast(scale); +} + +template +[[kernel]] void sdpa_blocked_causal_mask( + device T* scores [[buffer(0)]], + constant int& tile_rows [[buffer(1)]], + constant int& columns [[buffer(2)]], + uint2 elem [[thread_position_in_grid]]) { + int column = elem.x; + int score_row = elem.y; + int row = score_row % tile_rows; + if (column > row) { + size_t index = size_t(score_row) * columns + columns - tile_rows + column; + scores[index] = Limits::finite_min; + } +} diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h index 3f40dbd7..1d312038 100644 --- a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h +++ b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h @@ -61,6 +61,7 @@ template const int q_batch_head_idx = tid.x; const int q_seq_idx = tid.y; const int kv_head_idx = q_batch_head_idx / gqa_factor; + const int kv_seq_idx = simd_gid; const int o_offset = q_batch_head_idx * tpg.y + q_seq_idx; const int q_offset = query_transposed ? tpg.x * q_seq_idx + q_batch_head_idx : o_offset; @@ -71,11 +72,11 @@ template simd_lid * v_per_thread; if (bool_mask) { bmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } if (float_mask) { fmask += q_batch_head_idx * mask_head_stride + - simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; + kv_seq_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride; } out += o_offset * V + simd_gid * v_per_thread; @@ -176,7 +177,7 @@ template } } -template +template [[kernel]] void sdpa_vector_2pass_1( const device T* queries [[buffer(0)]], const device T* keys [[buffer(1)]], @@ -259,8 +260,73 @@ template sum_exp_score = 1; } + const int stride_k_tok = blocks * int(k_seq_stride); + const int stride_v_tok = blocks * int(v_seq_stride); + const int stride_mask_tok = + (bool_mask || float_mask) ? blocks * mask_kv_seq_stride : 0; + // For each key - for (int i = block_idx; i < N; i += blocks) { + int i = block_idx; + for (; i + (TK - 1) * blocks < N; i += blocks * TK) { + // Prefetch: issue all TK K/V loads first so the HW can overlap them. + T k_reg[TK][qk_per_thread]; + T v_reg[TK][v_per_thread]; + for (int t = 0; t < TK; t++) { + const device T* kp = keys + t * stride_k_tok; + for (int j = 0; j < qk_per_thread; j++) { + k_reg[t][j] = kp[j]; + } + } + for (int t = 0; t < TK; t++) { + const device T* vp = values + t * stride_v_tok; + for (int j = 0; j < v_per_thread; j++) { + v_reg[t][j] = vp[j]; + } + } + + for (int t = 0; t < TK; t++) { + int ii = i + t * blocks; + bool use_key = true; + if (do_causal) { + use_key = ii <= (N - q_seq_len + int(q_seq_idx)); + } else if (bool_mask) { + use_key = bmask[t * stride_mask_tok]; + } else if (float_mask) { + use_key = (fmask[t * stride_mask_tok] >= Limits::finite_min); + } + if (use_key) { + U score = 0; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * static_cast(k_reg[t][j]); + } + score = simd_sum(score); + if (float_mask) { + score += static_cast(fmask[t * stride_mask_tok]); + } + + U new_max = max(max_score, score); + U factor = fast::exp(max_score - new_max); + U exp_score = fast::exp(score - new_max); + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * static_cast(v_reg[t][j]); + } + } + } + + keys += TK * stride_k_tok; + values += TK * stride_v_tok; + if (bool_mask) { + bmask += TK * stride_mask_tok; + } + if (float_mask) { + fmask += TK * stride_mask_tok; + } + } + + // Tail (fewer than TK remaining keys) + for (; i < N; i += blocks) { bool use_key = true; if (do_causal) { use_key = i <= (N - q_seq_len + int(q_seq_idx)); @@ -272,8 +338,8 @@ template if (use_key) { // Compute the i-th score U score = 0; - for (int i = 0; i < qk_per_thread; i++) { - score += q[i] * keys[i]; + for (int j = 0; j < qk_per_thread; j++) { + score += q[j] * keys[j]; } score = simd_sum(score); @@ -290,19 +356,19 @@ template sum_exp_score = sum_exp_score * factor + exp_score; // Update the output accumulator - for (int i = 0; i < v_per_thread; i++) { - o[i] = o[i] * factor + exp_score * values[i]; + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + exp_score * values[j]; } } // Move the pointers to the next kv - keys += blocks * int(k_seq_stride); - values += blocks * int(v_seq_stride); + keys += stride_k_tok; + values += stride_v_tok; if (bool_mask) { - bmask += blocks * mask_kv_seq_stride; + bmask += stride_mask_tok; } if (float_mask) { - fmask += blocks * mask_kv_seq_stride; + fmask += stride_mask_tok; } } @@ -357,6 +423,7 @@ template const int num_kv_heads = tpg.x; const int num_q_heads = num_kv_heads * G; const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx; const int chunk = (N + blocks - 1) / blocks; const int kstart = block_idx * chunk; @@ -365,9 +432,9 @@ template const int s0 = kstart + cchunk * sub; const int s1 = min(kend, s0 + sub); - const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + - simd_lid * qk_per_thread; - const device T* vp = values + kv_head_idx * v_head_stride + + const device T* kp = keys + kv_batch_head_idx * k_head_stride + + s0 * k_seq_stride + simd_lid * qk_per_thread; + const device T* vp = values + kv_batch_head_idx * v_head_stride + s0 * v_seq_stride + simd_lid * v_per_thread; U q[HPT][qk_per_thread]; diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h index 1926aa3c..53a4017c 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "../../../steel/attn/attn.h" @@ -14,6 +14,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; struct MaxOp { template @@ -63,6 +64,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -76,6 +78,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -102,6 +105,12 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -114,13 +123,15 @@ template < constexpr short LDQ_tgp = BD + padQ; constexpr short LDK_tgp = BK + padK; - constexpr short LDV_tgp = BD + padV; + constexpr short LDV_tgp = BV + padV; constexpr short tgp_mem_0 = (BK + padK) * (BD); - constexpr short tgp_mem_1 = BK * (BD + padV); + constexpr short tgp_mem_1 = BK * (BV + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = (WN == 2) && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -150,7 +161,7 @@ template < using VBlockLoader = BlockLoaderT< /* typename T = */ T, /* short BROWS = */ BK, - /* short BCOLS = */ BD, + /* short BCOLS = */ BV, /* short kDstStrRow = */ LDV_tgp, /* short kDstStrCol = */ 1, /* short reduction_dim = */ 0, @@ -168,8 +179,15 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; + + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert( + BD % WN == 0 && BV % WN == 0, + "The head dims must split evenly across WN"); - constexpr int kNWarps = WM * WN; + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -180,14 +198,19 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + constexpr int TVh = BV / (kFragSize * WN); + constexpr int BVh = BV / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -195,22 +218,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BVh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -262,17 +308,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -373,7 +453,7 @@ template < // Load V blocks if (!align_K && kb == (params->NK_aligned)) { - loader_v.load_safe(short2(BD, params->kL_rem)); + loader_v.load_safe(short2(BV, params->kL_rem)); } else { loader_v.load_unsafe(); } @@ -425,7 +505,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TVh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -460,11 +540,25 @@ template < Otile.template row_bin_op(sum_score); threadgroup_barrier(mem_flags::mem_none); + // Store the logsumexp for the backward pass + if (save_lse && sn == 0) { + using stile_t = decltype(Stile); + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + sm + i * stile_t::kFragRows; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = M_LN2_F * (max_score[i] + metal::log2(sum_score[i])); + } + } + } + // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BVh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; @@ -473,4 +567,4 @@ template < } else { Otile.template store(O, params->O_strides[2]); } -} +} \ No newline at end of file diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal index 9efa26fe..7d0fe145 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal @@ -1,31 +1,36 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. // clang-format off #include "../../../utils.h" #include "../../../steel/attn/kernels/steel_attention.h" -#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ - instantiate_kernel( \ - "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ - "_wm" #wm "_wn" #wn "_mask" #mname, \ - attention, dtype, bq, bk, bd, wm, wn, mtype, float) +#define instantiate_attn(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention, dtype, bq, bk, bd, bv, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) +#define instantiate_attn_helper(iname, itype, bq, bk, bd, bv, wm, wn) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, iname, itype) \ + instantiate_attn(iname, itype, bq, bk, bd, bv, wm, wn, bool_, bool) -#define instantiate_attn_mask_helper(iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, bool_, bool) +#define instantiate_attn_shapes_helper(iname, itype) \ + instantiate_attn_helper(iname, itype, 32, 16, 192, 192, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 16, 128, 128, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 96, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 64, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 80, 80, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 72, 72, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 64, 64, 4, 1) -instantiate_attn_mask_helper(float16, half); -instantiate_attn_mask_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float16, half); +instantiate_attn_shapes_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float32, float); -instantiate_attn_mask_helper(float32, float); +instantiate_attn_helper(float16, half, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(float16, half, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 16, 256, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 32, 256, 256, 4, 2); +instantiate_attn_helper(float32, float, 32, 16, 256, 256, 4, 2); // clang-format on diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h index 5010f1de..43dc2db8 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "../../../steel/attn/nax.h" #include "../../../steel/attn/params.h" @@ -17,6 +17,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; template struct TransformScale { @@ -76,6 +77,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -89,6 +91,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -116,6 +119,13 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + // [B, H, qL], contiguous. + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -136,11 +146,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -421,8 +433,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } @@ -472,6 +484,16 @@ template < Otile.template row_bin_op(rcp); + if (save_lse && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + i * otile_t::kFragRowsJump + sm; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = max_score[i] * M_LN2_F + metal::log(sum_score[i]); + } + } + } + // Store results O += tm * int(params->O_strides[2]); @@ -489,15 +511,10 @@ template < // Head-dim split attention kernel /////////////////////////////////////////////////////////////////////////////// -// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup -// accumulator working set of attention_nax (TD output fragments plus the S -// fragments) is what gates tensor-unit throughput, so this kernel splits the -// head dim across the WN = 2 simdgroups of the second warp dimension: each -// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, -// halving its accumulator set. The pair exchanges its partial Q@K.T sums -// through threadgroup memory, then both simdgroups run softmax redundantly on -// the full S tile (the row statistics are cheap) and each accumulates P@V for -// its own half of Dv. +// Split wide heads across WN simdgroups to reduce the accumulator working set. +// Each group owns one D slice. The groups exchange their partial Q@K.T scores +// through threadgroup memory and reduce them in the same order. Each group then +// runs softmax and accumulates its P@V slice. // clang-format off template < @@ -518,6 +535,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* LSE [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -556,9 +574,9 @@ template < constexpr short kU = 16; // The WM simdgroups along the first warp dimension split the Q sequence; - // the WN simdgroups along the second split the head dim. The exchange - // below reduces exactly one peer, so WN is fixed at 2. - static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + // the WN simdgroups along the second split the head dim. + static_assert( + WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4"); constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, @@ -579,20 +597,20 @@ template < constexpr int BDh = BD / WN; static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); - static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + static_assert(TK % 2 == 0, "S fragments are computed pair by pair"); const short row_group = simd_group_id / WN; - const short d_half = simd_group_id % WN; + const short d_group = simd_group_id % WN; using otile_t = NAXTile; otile_t Otile; Otile.clear(); const short tm = kU * TQ * row_group; - Q += tm * int(params->Q_strides[2]) + d_half * BDh; - K += d_half * BDh; - V += d_half * BDh; - O += tm * int(params->O_strides[2]) + d_half * BDh; + Q += tm * int(params->Q_strides[2]) + d_group * BDh; + K += d_group * BDh; + V += d_group * BDh; + O += tm * int(params->O_strides[2]) + d_group * BDh; constexpr short kRowsPT = otile_t::kRowsPerThread; @@ -632,12 +650,11 @@ template < using stile_t = NAXTile; constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; - // One slot per (row group, half): a fragment pair in per-lane-linear - // layout. Both halves share the fragment-to-lane mapping, so the - // exchange needs no coordinate math. - threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + // One slot per (row group, D group) in per-lane-linear layout. All groups + // share the fragment-to-lane mapping. + threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32]; - // Keep the simdgroup's Q half resident in registers for the whole KV + // Keep the simdgroup's Q slice resident in registers for the whole KV // loop: TDh fragments of T are cheap next to the accumulators. NAXTile Qtiles[TDh]; STEEL_PRAGMA_UNROLL @@ -662,7 +679,7 @@ template < stile_t Stile; Stile.clear(); - // S = Q @ K.T, this half of D only, exchanged pair by pair. + // S = Q @ K.T for this D slice. STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { STEEL_PRAGMA_UNROLL @@ -686,26 +703,37 @@ template < Ktile.frag_at(1, 0), metal::true_type{}); } + } - // Exchange the partial pair and reduce. - threadgroup AccumType* slot = s_xchg[row_group][d_half]; - thread auto& s0 = Stile.frag_at(0, ik); - thread auto& s1 = Stile.frag_at(0, ik + 1); - const short base = short(simd_lane_id) * (2 * kEPF); + // Exchange all partial scores and reduce them in a fixed order. + threadgroup AccumType* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - slot[base + i] = s0[i]; - slot[base + kEPF + i] = s1[i]; + slot[base + ik * kEPF + i] = s[i]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - s0[i] += peer[base + i]; - s1[i] += peer[base + kEPF + i]; + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short peer_group = 1; peer_group < WN; peer_group++) { + const threadgroup AccumType* peer = s_xchg[row_group][peer_group]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += peer[base + ik * kEPF + i]; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); // Scale S STEEL_PRAGMA_UNROLL @@ -826,7 +854,7 @@ template < } } - // Do softmax (redundantly per half; the row statistics are cheap) + // Do softmax in each D group; the row statistics are cheap. metal::vec new_max; metal::vec factor; STEEL_PRAGMA_UNROLL @@ -854,7 +882,7 @@ template < simdgroup_barrier(mem_flags::mem_none); - // O = P @ V, this half of Dv only. + // O = P @ V for this D slice. STEEL_PRAGMA_UNROLL for (short id = 0; id < TDh; id += 2) { STEEL_PRAGMA_UNROLL @@ -895,6 +923,21 @@ template < rcp[i] = 1.f / sum_score[i]; } + if (save_lse && d_group == 0 && sn == 0) { + const int lse_row_base = int(tid.x) * BQ + tm; + const int lse_head_off = (int(tid.z) * params->H + int(tid.y)) * params->qL; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = lse_row_base + i * stile_t::kFragRowsJump + sm; + if (row < params->qL) { + LSE[lse_head_off + row] = (sum_score[i] == 0) + ? -INFINITY + : float(M_LN2_F * (max_score[i] + metal::log2(sum_score[i]))); + } + } + } + Otile.template row_bin_op(rcp); if (!align_Q && is_last_q) { diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h index f559596b..f948ec06 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h @@ -46,7 +46,7 @@ implicit_gemm_conv_2d( using loader_a_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DInputBlockLoaderSmallChannels< @@ -84,7 +84,7 @@ implicit_gemm_conv_2d( // Weight loader using loader_b_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DWeightBlockLoaderSmallChannels< diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h index d0797ebe..de623ee4 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h @@ -186,7 +186,7 @@ implicit_gemm_conv_2d_general( if (offset_n >= gemm_params->N) return; - short diff = gemm_params->N - offset_n; + int diff = gemm_params->N - offset_n; STEEL_PRAGMA_UNROLL for (int i = 0; i < mma_t::TM; i++) { diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h index 1284efe4..9c134007 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h @@ -47,8 +47,8 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -197,8 +197,8 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params; const constant ImplicitGemmConv2DParams* gemm_params; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -384,6 +384,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -405,6 +406,7 @@ struct Conv2DWeightBlockLoader { weight_hw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -419,7 +421,7 @@ struct Conv2DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; @@ -481,9 +483,9 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; short kdil_d; short kdil_h; @@ -665,9 +667,9 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params; const constant ImplicitGemmConv3DParams* gemm_params; - short weight_d; - short weight_h; - short weight_w; + int weight_d; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -887,6 +889,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -908,6 +911,7 @@ struct Conv3DWeightBlockLoader { weight_dhw(0), weight_step(params->C / params->groups), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -922,7 +926,7 @@ struct Conv3DWeightBlockLoader { } } else { for (short i = 0; i < BN; i += TROWS) { - if ((read_n + i) < params->O) { + if ((read_n + i) < N) { STEEL_PRAGMA_UNROLL for (short j = 0; j < vec_size; j++) { dst[i * dst_ld + j] = src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h index a90e9c88..a2016fb5 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h @@ -71,6 +71,7 @@ struct Conv2DInputBlockLoaderSmallChannels { // Rows / strided reads within the block STEEL_CONST short n_rows = BROWS / TROWS; + static_assert(n_rows > 0, "BM must be at least tgp_size * vec_size / BK"); // Thread location indices const short thread_idx; @@ -233,6 +234,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -253,6 +255,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { params(params_), weight_hw(thread_idx % TCOLS), read_n(offsets.y + bi), + N(gemm_params_->N), do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ @@ -260,7 +263,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { if (bi >= BROWS || bj >= BCOLS) return; - if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -289,7 +292,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } } else { for (short i = 0; i < BROWS; i += TROWS) { - if (((read_n + i) < params->O)) { + if (((read_n + i) < N)) { STEEL_PRAGMA_UNROLL for (short j = 0; j < n_channels; j++) { dst[i * dst_ld + j] = curr_src[i * src_ld + j]; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h index b44ff341..1f4c6602 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h @@ -48,8 +48,8 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -248,8 +248,8 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const int start_row; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h index 4c055e69..54e45d1e 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather.h @@ -20,9 +20,10 @@ template < [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { @@ -44,8 +45,19 @@ template < using loader_b_t = typename gemm_kernel::loader_b_t; using mma_t = typename gemm_kernel::mma_t; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } @@ -53,188 +65,127 @@ template < threadgroup T As[gemm_kernel::tgp_mem_size_a]; threadgroup T Bs[gemm_kernel::tgp_mem_size_b]; - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; const size_t c_row_long = size_t(c_row); const size_t c_col_long = size_t(c_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row)); const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col)); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[c_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (rhs_indices[c_row + n] != index) { - offset_next = n; - index_next = rhs_indices[c_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); + // Prepare threadgroup loading operations + thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); + thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id); + + // Prepare iterations + const int gemm_k_iterations = params->gemm_k_iterations_aligned; - // Prepare threadgroup loading operations - thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); - thread loader_b_t loader_b( - B + index * params->batch_stride_b, - params->ldb, + // Do unaligned K iterations first + if (!align_K) { + const int k_last = params->gemm_k_iterations_aligned * BK; + const int k_remain = params->K - k_last; + const size_t k_jump_a = + transpose_a ? params->lda * size_t(k_last) : size_t(k_last); + const size_t k_jump_b = + transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); + + // Move loader source ahead to end + loader_a.src += k_jump_a; + loader_b.src += k_jump_b; + + // Load tile + const short2 tile_dims_A = + transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); + const short2 tile_dims_B = + transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); + + loader_a.load_safe(tile_dims_A); + loader_b.load_safe(tile_dims_B); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Do matmul + mma_op.mma(As, Bs); + + // Reset source back to start + loader_a.src -= k_jump_a; + loader_b.src -= k_jump_b; + } + + const short lbk = 0; + + // Tile aligned don't check + if ((tgp_bm == BM) && (align_N || tgp_bn == BN)) { + gemm_kernel::gemm_loop( + As, Bs, - simd_group_id, - simd_lane_id); - - // Prepare iterations - const int gemm_k_iterations = params->gemm_k_iterations_aligned; - - // Do unaligned K iterations first - if (!align_K) { - const int k_last = params->gemm_k_iterations_aligned * BK; - const int k_remain = params->K - k_last; - const size_t k_jump_a = - transpose_a ? params->lda * size_t(k_last) : size_t(k_last); - const size_t k_jump_b = - transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); - - // Move loader source ahead to end - loader_a.src += k_jump_a; - loader_b.src += k_jump_b; - - // Load tile - const short2 tile_dims_A = - transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); - const short2 tile_dims_B = - transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); - - loader_a.load_safe(tile_dims_A); - loader_b.load_safe(tile_dims_B); + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - threadgroup_barrier(mem_flags::mem_threadgroup); + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Do matmul - mma_op.mma(As, Bs); + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Reset source back to start - loader_a.src -= k_jump_a; - loader_b.src -= k_jump_b; - } + // Nothing aligned so check both rows and cols + else { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Matrix level aligned never check - if (align_M && align_N) { - for (int k = 0; k < gemm_k_iterations; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Load elements into threadgroup - loader_a.load_unsafe(); - loader_b.load_unsafe(); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Multiply and accumulate threadgroup elements - mma_op.mma(As, Bs); - - // Prepare for next iteration - loader_a.next(); - loader_b.next(); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } else { - const short lbk = 0; - - // Tile aligned don't check - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } - - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - - // Nothing aligned so check both rows and cols - else { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - } + // Store results to device memory + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + mma_op.store_result(C, params->ldd); + } else { + mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h index 7fec2706..afd645ff 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h @@ -20,9 +20,11 @@ template < gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], + uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { constexpr short SM = BM / WM; @@ -31,106 +33,66 @@ gather_mm_rhs_nax( constexpr short TM = SM / 16; constexpr short TN = SN / 16; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; - const size_t c_row_long = size_t(c_row); - const size_t c_col_long = size_t(c_col); + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + const short sgp_sm = short(clamp(tgp_bm - tm, 0, int(SM))); + const short sgp_sn = + align_N ? SN : short(clamp(params->N - (c_col + tn), 0, int(SN))); + const size_t c_row_long = size_t(c_row + tm); + const size_t c_col_long = size_t(c_col + tn); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - rhs_indices += c_row; - - const short tm = SM * (simd_group_id / WN); - const short tn = SN * (simd_group_id % WN); - - const int sgp_sm_int = - align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); - const short sgp_sm = short(sgp_sm_int); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); - const int sgp_sn_int = - align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); - const short sgp_sn = short(sgp_sn_int); - const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); + dispatch_bool(align_K, [&](auto kAlignedK) { + dispatch_bool(sgp_sm == SM, [&](auto kAlignedM) { + dispatch_bool(sgp_sn == SN, [&](auto kAlignedN) { + NAXTile Ctile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + params->K, + params->gemm_k_iterations_aligned, + sgp_sm, + sgp_sn); - A += transpose_a ? tm : (tm * params->lda); - B += transpose_b ? (tn * params->ldb) : tn; - C += tm * params->ldd + tn; - rhs_indices += tm; - - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[0]; - short offset_next = 0; - int n = 0; - while (n < sgp_sm) { - n++; - offset = offset_next; - index = index_next; - offset_next = sgp_sm; - for (; n < sgp_sm; n++) { - if (rhs_indices[n] != index) { - offset_next = n; - index_next = rhs_indices[n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - NAXTile Ctile; - - dispatch_bool(align_K, [&](auto kAlignedK) { - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K - T, - SM, - SN, - SK, - BK, - transpose_a, - transpose_b, - kAlignedM.value, - kAlignedN.value, - kAlignedK.value, - AccumType>; - Ctile = do_gemm( - A, - B + index * params->batch_stride_b, - params->lda, - params->ldb, - params->K, - params->gemm_k_iterations_aligned, - sgp_sm, - sgp_sn); - - if constexpr (kAlignedN.value) { - if (offset_next - offset == SM) { - Ctile.store(C, int(params->ldd)); - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(SN, offset_next)); - } - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(sgp_sn, offset_next)); - } - }); + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.store(C, int(params->ldd)); + } else { + Ctile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } }); }); - } + }); } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h index 5a43e223..6814ef50 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented.h @@ -79,8 +79,8 @@ template < k_start = segments[tid.z]; k_end = segments[tid.z + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid.z * params->batch_stride_d; // Prepare threadgroup mma operation diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h index 9cfece5c..5731cb8c 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h @@ -51,8 +51,8 @@ void segmented_mm_nax( k_start = segments[tid_s]; k_end = segments[tid_s + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid_s * params->batch_stride_d; constexpr short SM = BM / WM; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h new file mode 100644 index 00000000..28229a03 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_thin_nax.h @@ -0,0 +1,223 @@ +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} diff --git a/Source/Cmlx/mlx-generated/metal/unary_ops.h b/Source/Cmlx/mlx-generated/metal/unary_ops.h index a2dc8042..5909350a 100644 --- a/Source/Cmlx/mlx-generated/metal/unary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/unary_ops.h @@ -35,7 +35,7 @@ struct Abs { return x; }; complex64_t operator()(complex64_t x) thread { - return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; + return {hypot(x.real, x.imag), 0}; }; }; @@ -308,7 +308,7 @@ struct Round { struct Sigmoid { template T operator()(T x) thread { - auto y = 1 / (1 + metal::exp(metal::abs(x))); + auto y = 1 / (1 + metal::precise::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -325,8 +325,8 @@ struct Sign { if (x == complex64_t(0)) { return x; } - return x / - (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag); + auto r = hypot(x.real, x.imag); + return {x.real / r, x.imag / r}; }; }; @@ -378,7 +378,7 @@ struct Sqrt { auto b_abs = metal::precise::sqrt((r - x.real) / 2.0); auto b = metal::copysign(b_abs, x.imag); return {a, b}; - } + }; }; struct Rsqrt { diff --git a/Source/Cmlx/mlx-generated/metal/utils.h b/Source/Cmlx/mlx-generated/metal/utils.h index cbf5798a..7c82bf3c 100644 --- a/Source/Cmlx/mlx-generated/metal/utils.h +++ b/Source/Cmlx/mlx-generated/metal/utils.h @@ -353,6 +353,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -470,3 +493,50 @@ template <> inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } + +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 583bba8a..22009b9a 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -41,15 +41,30 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -59,8 +74,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -74,7 +89,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -84,8 +99,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -99,7 +114,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -120,15 +135,30 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1]; + x_thread[i + 2] = x[i + 2]; + x_thread[i + 3] = x[i + 3]; + x_thread[i + 4] = x[i + 4]; + x_thread[i + 5] = x[i + 5]; + x_thread[i + 6] = x[i + 6]; + x_thread[i + 7] = x[i + 7]; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -138,8 +168,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -154,7 +184,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -164,8 +194,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -179,7 +209,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -209,13 +239,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -311,13 +355,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + uint8_t wb = w[i]; + accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); + accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); + accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); + accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); + accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); + accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); + accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); + accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -408,11 +466,25 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + uint8_t wb = w[i]; + result[8 * i] += x * (select(U(0), scale, bool(wb & 0x01)) + bias); + result[8 * i + 1] += x * (select(U(0), scale, bool(wb & 0x02)) + bias); + result[8 * i + 2] += x * (select(U(0), scale, bool(wb & 0x04)) + bias); + result[8 * i + 3] += x * (select(U(0), scale, bool(wb & 0x08)) + bias); + result[8 * i + 4] += x * (select(U(0), scale, bool(wb & 0x10)) + bias); + result[8 * i + 5] += x * (select(U(0), scale, bool(wb & 0x20)) + bias); + result[8 * i + 6] += x * (select(U(0), scale, bool(wb & 0x40)) + bias); + result[8 * i + 7] += x * (select(U(0), scale, bool(wb & 0x80)) + bias); + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -498,14 +570,36 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { template inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -598,9 +692,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -609,6 +703,12 @@ struct QuantizedBlockLoader { (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0; + + static_assert( + group_size % (n_reads * pack_factor) == 0, + "The group size must be a multiple of the columns read per thread."); + const int src_ld; const int tile_stride; short group_step_cnt; @@ -664,14 +764,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -767,7 +860,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -779,7 +872,8 @@ METAL_FUNC void qmv_fast_impl( uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { - constexpr int packs_per_thread = bits == 2 ? 1 : 2; + constexpr int packs_per_thread = + bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor(); @@ -801,19 +895,37 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; x += tid.x * in_vec_size + simd_lid * values_per_thread; y += tid.x * out_vec_size + out_row; - for (int k = 0; k < in_vec_size; k += block_size) { + const int aligned_end = (in_vec_size / block_size) * block_size; + + for (int k = 0; k < aligned_end; k += block_size) { U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -826,9 +938,37 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } + // Tail block: handles K that is a multiple of 512 (dispatch gate) but not + // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for + // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; + // their result[row] stays at 0 and contributes 0 to the simd_sum below. + if (aligned_end < in_vec_size) { + bool in_bounds = + (aligned_end + static_cast(simd_lid) * values_per_thread) < + in_vec_size; + if (in_bounds) { + U sum = load_vector(x, x_thread); + + for (int row = 0; row < results_per_simdgroup; row++) { + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; + + U s = sl[0]; + U b = bl[0]; + result[row] += + qdot(wl, x_thread, s, b, sum); + } + } + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1207,6 +1347,7 @@ template < const int group_size, const int bits, const bool aligned_N, + typename U = T, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1215,7 +1356,7 @@ METAL_FUNC void qmm_t_impl( const device T* scales, const device T* biases, const device T* x, - device T* y, + device U* y, threadgroup T* Xs, threadgroup T* Ws, const constant int& K, @@ -1240,7 +1381,7 @@ METAL_FUNC void qmm_t_impl( // Instantiate the appropriate BlockMMA and Loader using mma_t = mlx::steel:: - BlockMMA; + BlockMMA; using loader_x_t = mlx::steel::BlockLoader; using loader_w_t = QuantizedBlockLoader< @@ -1616,6 +1757,7 @@ template < int group_size, int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( @@ -1656,7 +1798,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1674,6 +1816,7 @@ template < int group_size, const int bits, bool batched, + bool partial_rows = false, bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_qmv( @@ -1910,6 +2053,7 @@ template < const int bits, const bool aligned_N, const bool batched, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -1959,7 +2103,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -1990,7 +2134,7 @@ template < const device T* scales [[buffer(1)]], const device T* biases [[buffer(2)]], const device T* x [[buffer(3)]], - device T* y [[buffer(4)]], + device float* y [[buffer(4)]], const constant int& K [[buffer(5)]], const constant int& N [[buffer(6)]], const constant int& M [[buffer(7)]], @@ -2018,7 +2162,7 @@ template < biases += k_start / group_size; y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2289,6 +2433,7 @@ template < const int group_size, const int bits, const bool aligned_N, + const bool has_global_scale = false, const int BM = 32, const int BK = 32, const int BN = 32> @@ -2348,7 +2493,7 @@ template < s_strides, b_strides, tid); - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2450,11 +2595,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -2498,13 +2644,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -2521,113 +2672,61 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); - - // Prepare threadgroup loading operations - thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - // Matrices are all aligned check nothing - if (align_M && align_N) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } else { - // Tile aligned so check outside of the hot loop - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } + // Prepare threadgroup loading operations + thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id); + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(y, N); - } else { - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } - } + // Tile aligned so check outside of the hot loop + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result(y, N); + } - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(BN, offset_next)); - } + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(BN, tgp_bm)); + } - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); + } + mma_op.store_result_safe(y, N, short2(tgp_bn, BM)); + } - // Nothing aligned so check both rows and cols - else { - gemm_loop_unaligned( - Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - gemm_loop_finalize( - Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); - } - mma_op.store_result_slice( - y, N, short2(0, offset), short2(tgp_bn, offset_next)); - } + // Nothing aligned so check both rows and cols + else { + gemm_loop_unaligned( + Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK); + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w); } + mma_op.store_result_safe(y, N, short2(tgp_bn, tgp_bm)); } } @@ -2675,14 +2774,23 @@ template w_min = simd_min(w_min); w_max = simd_max(w_max); - float scale = max((w_max - w_min) / n_bins, eps); - bool side = abs(w_min) > abs(w_max); - scale = side ? scale : -scale; - float edge = side ? w_min : w_max; - float q0 = round(edge / scale); - bool at_zero = q0 == 0.0f; - scale = at_zero ? scale : edge / q0; - float bias = at_zero ? 0 : edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = max(w_max - w_min, eps); + bias = w_min; + } else { + scale = max((w_max - w_min) / n_bins, eps); + bool side = abs(w_min) > abs(w_max); + scale = side ? scale : -scale; + float edge = side ? w_min : w_max; + float q0 = round(edge / scale); + bool at_zero = q0 == 0.0f; + scale = at_zero ? scale : edge / q0; + bias = at_zero ? 0 : edge; + } // Write out the scales and biases size_t gindex = in_index / group_size; @@ -2786,7 +2894,9 @@ template #pragma clang loop unroll(full) for (int i = 0; i < pack_factor; i++) { uint8_t d; - if (bits == 2) { + if (bits == 1) { + d = (val >> i) & 0x01; + } else if (bits == 2) { d = (val >> (bits * i)) & 0x03; } else if (bits == 4) { d = (val >> (bits * i)) & 0x0f; diff --git a/Source/Cmlx/mlx-generated/quantized_nax.cpp b/Source/Cmlx/mlx-generated/quantized_nax.cpp index 123a4a37..8f5d2297 100644 --- a/Source/Cmlx/mlx-generated/quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/quantized_nax.cpp @@ -44,15 +44,30 @@ inline constexpr short get_bytes_per_pack() { template inline U load_vector(const device T* x, thread U* x_thread) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < values_per_thread; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -62,8 +77,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -77,7 +92,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -87,8 +102,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -102,7 +117,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -123,15 +138,30 @@ inline U load_vector(const device T* x, thread U* x_thread) { template inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U sum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < N; i += 8) { + sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + + x[i + 6] + x[i + 7]; + x_thread[i] = x[i]; + x_thread[i + 1] = x[i + 1] / 2.0f; + x_thread[i + 2] = x[i + 2] / 4.0f; + x_thread[i + 3] = x[i + 3] / 8.0f; + x_thread[i + 4] = x[i + 4] / 16.0f; + x_thread[i + 5] = x[i + 5] / 32.0f; + x_thread[i + 6] = x[i + 6] / 64.0f; + x_thread[i + 7] = x[i + 7] / 128.0f; + } + } + + else if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -141,8 +171,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -157,7 +187,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -167,8 +197,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -182,7 +212,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -212,13 +242,27 @@ inline U qdot( U bias, U sum) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (values_per_thread / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (values_per_thread / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -314,13 +358,27 @@ inline U qdot_safe( U sum, int N) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); U accum = 0; - if (bits == 2) { + if (bits == 1) { + for (int i = 0; i < (N / 8); i++) { + accum += + (x_thread[8 * i] * (w[i] & 0x01) + + x_thread[8 * i + 1] * (w[i] & 0x02) + + x_thread[8 * i + 2] * (w[i] & 0x04) + + x_thread[8 * i + 3] * (w[i] & 0x08) + + x_thread[8 * i + 4] * (w[i] & 0x10) + + x_thread[8 * i + 5] * (w[i] & 0x20) + + x_thread[8 * i + 6] * (w[i] & 0x40) + + x_thread[8 * i + 7] * (w[i] & 0x80)); + } + } + + else if (bits == 2) { for (int i = 0; i < (N / 4); i++) { accum += (x_thread[4 * i] * (w[i] & 0x03) + @@ -411,11 +469,33 @@ template inline void qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); + + if (bits == 1) { + U s[8] = { + scale, + scale / 2.0f, + scale / 4.0f, + scale / 8.0f, + scale / 16.0f, + scale / 32.0f, + scale / 64.0f, + scale / 128.0f}; + for (int i = 0; i < (values_per_thread / 8); i++) { + result[8 * i] += x * (s[0] * (w[i] & 0x01) + bias); + result[8 * i + 1] += x * (s[1] * (w[i] & 0x02) + bias); + result[8 * i + 2] += x * (s[2] * (w[i] & 0x04) + bias); + result[8 * i + 3] += x * (s[3] * (w[i] & 0x08) + bias); + result[8 * i + 4] += x * (s[4] * (w[i] & 0x10) + bias); + result[8 * i + 5] += x * (s[5] * (w[i] & 0x20) + bias); + result[8 * i + 6] += x * (s[6] * (w[i] & 0x40) + bias); + result[8 * i + 7] += x * (s[7] * (w[i] & 0x80) + bias); + } + } - if (bits == 2) { + else if (bits == 2) { U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f}; for (int i = 0; i < (values_per_thread / 4); i++) { result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias); @@ -500,14 +580,36 @@ template inline void dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); const float s = float(scale); const float b = float(bias); - if (bits == 2) { + if (bits == 1) { + float sc[8] = { + s, + s / 2.0f, + s / 4.0f, + s / 8.0f, + s / 16.0f, + s / 32.0f, + s / 64.0f, + s / 128.0f}; + for (int i = 0; i < (N / 8); i++) { + w_local[8 * i] = static_cast(sc[0] * (w[i] & 0x01) + b); + w_local[8 * i + 1] = static_cast(sc[1] * (w[i] & 0x02) + b); + w_local[8 * i + 2] = static_cast(sc[2] * (w[i] & 0x04) + b); + w_local[8 * i + 3] = static_cast(sc[3] * (w[i] & 0x08) + b); + w_local[8 * i + 4] = static_cast(sc[4] * (w[i] & 0x10) + b); + w_local[8 * i + 5] = static_cast(sc[5] * (w[i] & 0x20) + b); + w_local[8 * i + 6] = static_cast(sc[6] * (w[i] & 0x40) + b); + w_local[8 * i + 7] = static_cast(sc[7] * (w[i] & 0x80) + b); + } + } + + else if (bits == 2) { float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); @@ -600,9 +702,9 @@ struct QuantizedBlockLoader { group_size % BCOLS == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -666,14 +768,7 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -735,9 +830,9 @@ struct QuantizedBlockLoader< BCOLS % group_size == 0, "The group size should be divisible by the columns"); static_assert( - bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || - bits == 8, - "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + bits == 1 || bits == 2 || bits == 3 || bits == 4 || bits == 5 || + bits == 6 || bits == 8, + "Template undefined for bits not in {1, 2, 3, 4, 5, 6, 8}"); MLX_MTL_CONST short pack_factor = get_pack_factor(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); @@ -806,14 +901,7 @@ struct QuantizedBlockLoader< return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { - for (int i = 0; i < n_reads * pack_factor; i++) { - dst[i] = T(0); - } - return; - } - - if (reduction_dim == 0 && bi >= src_tile_dim.y) { + if (bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); } @@ -1228,7 +1316,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1347,7 +1436,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + const bool has_global_scale = false> [[kernel]] void affine_gather_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1490,11 +1580,12 @@ template < const device uint32_t* w [[buffer(1)]], const device T* scales [[buffer(2)]], const device T* biases [[buffer(3)]], - const device uint32_t* indices [[buffer(4)]], + const device int32_t* offsets [[buffer(4)]], device T* y [[buffer(5)]], const constant int& M [[buffer(6)]], const constant int& N [[buffer(7)]], const constant int& K [[buffer(8)]], + const constant int& num_groups [[buffer(9)]], uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { @@ -1523,13 +1614,18 @@ template < const int K_it = K / BK; const size_t stride_w = transpose ? N * K_w : K * N_w; const size_t stride_s = transpose ? N * K_g : K * N_g; - const int y_row = tid.y * BM; + int y_row; + int group; + short tgp_bm; + if (!schedule_row_tile( + offsets, num_groups, M, tid.y, simd_lane_id, y_row, group, tgp_bm)) { + return; + } const int y_col = tid.x * BN; const size_t y_row_long = size_t(y_row); const size_t y_col_long = size_t(y_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, M - y_row)); const short tgp_bn = align_N ? BN : short(min(BN, N - y_col)); // Calculate the final tiles in the case that K is not aligned @@ -1556,12 +1652,11 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); + const short sgp_sm = short(clamp(int(tgp_bm) - tm, 0, int(SM))); const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + const bool rows_in_bounds = y_row + tm + SM <= M; const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); constexpr short BR = transpose ? TN : TK; @@ -1569,147 +1664,121 @@ template < using AccumType = float; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = indices[y_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (indices[y_row + n] != index) { - offset_next = n; - index_next = indices[y_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - const bool sg_active = m_hi_lim > m_lo_lim; - - NAXTile Dtile; - Dtile.clear(); - - const device T* xn = x + tm * K; - - // Prepare threadgroup loading operations - thread loader_w_t loader_w( - wl + index * stride_w, - scales + index * stride_s, - biases + index * stride_s, - transpose ? K : N, - Ws, - simd_group_id, - simd_lane_id); - - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { - for (int k = 0; k < K_it; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedN.value) { - loader_w.load_unsafe(); - } else { - loader_w.load_safe( - transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); - } + const bool sg_active = sgp_sm > 0; - threadgroup_barrier(mem_flags::mem_threadgroup); + NAXTile Dtile; + Dtile.clear(); - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + const device T* xn = x + tm * K; + + // Prepare threadgroup loading operations + thread loader_w_t loader_w( + wl + group * stride_w, + scales + group * stride_s, + biases + group * stride_s, + transpose ? K : N, + Ws, + simd_group_id, + simd_lane_id); + + dispatch_bool(rows_in_bounds, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) { + for (int k = 0; k < K_it; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (kAlignedN.value) { + loader_w.load_unsafe(); + } else { + loader_w.load_safe( + transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK)); + } - volatile int compiler_barrier; + threadgroup_barrier(mem_flags::mem_threadgroup); - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } - } - xn += BK; - loader_w.next(); - } + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - if (!align_K) { - threadgroup_barrier(mem_flags::mem_threadgroup); - loader_w.load_safe(tile_w); - threadgroup_barrier(mem_flags::mem_threadgroup); + (void)compiler_barrier; + } + } - STEEL_PRAGMA_NO_UNROLL - for (int kk1 = 0; kk1 < BK; kk1 += SK) { - if (sg_active) { - NAXTile Atile; - NAXTile Btile; + xn += BK; + loader_w.next(); + } - volatile int compiler_barrier; + if (!align_K) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_w.load_safe(tile_w); + threadgroup_barrier(mem_flags::mem_threadgroup); - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + volatile int compiler_barrier; - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + const short psk = min(int(SK), max(0, (k_remain - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - (void)compiler_barrier; + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } } + } - threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup_barrier(mem_flags::mem_threadgroup); - // Store results to device memory - if constexpr (kAlignedN.value) { - if (m_lo_lim == 0 && m_hi_lim == SM) { - Dtile.store(y + tm * N + tn, N); - } else { - Dtile.store_slice( - y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim)); - } + // Store results to device memory + if constexpr (kAlignedN.value) { + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); } else { Dtile.store_slice( - y + tm * N + tn, - N, - short2(0, m_lo_lim), - short2(sgp_sn, m_hi_lim)); + y + tm * N + tn, N, short2(0, 0), short2(SN, sgp_sm)); } - }); + } else { + Dtile.store_slice( + y + tm * N + tn, N, short2(0, 0), short2(sgp_sn, sgp_sm)); + } }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/reduce.cpp b/Source/Cmlx/mlx-generated/reduce.cpp index cac2d8de..3cb5e0ca 100644 --- a/Source/Cmlx/mlx-generated/reduce.cpp +++ b/Source/Cmlx/mlx-generated/reduce.cpp @@ -202,7 +202,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); @@ -257,6 +257,7 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], + const constant int64_t& column_offset [[buffer(11)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -278,7 +279,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT out_idx = gid.y + gsize.y * IdxT(gid.z); @@ -328,7 +329,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += out_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -396,7 +397,8 @@ template < const constant int64_t* reduce_strides [[buffer(8)]], const constant int& reduce_ndim [[buffer(9)]], const constant size_t& non_col_reductions [[buffer(10)]], - const constant size_t& out_size [[buffer(11)]], + const constant int64_t& column_offset [[buffer(11)]], + const constant size_t& out_size [[buffer(12)]], uint3 gid [[threadgroup_position_in_grid]], uint3 gsize [[threadgroups_per_grid]], uint simd_lane_id [[thread_index_in_simdgroup]], @@ -421,7 +423,7 @@ template < short lid = simd_group_id * simd_size + simd_lane_id; short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks); - IdxT column = BN * gid.x + offset.x; + IdxT column = IdxT(BN) * gid.x + IdxT(column_offset) + offset.x; bool safe = column + n_reads <= reduction_stride; IdxT full_idx = gid.y + gsize.y * IdxT(gid.z); @@ -468,7 +470,7 @@ template < // Write the output. if (simd_lane_id == 0) { - IdxT out_column = BN * gid.x + out_offset.x; + IdxT out_column = IdxT(BN) * gid.x + IdxT(column_offset) + out_offset.x; out += full_idx * IdxT(reduction_stride) + out_column; if (out_column + n_outputs <= reduction_stride) { for (int i = 0; i < n_outputs; i++) { @@ -492,8 +494,9 @@ template < template [[kernel]] void init_reduce( device T* out [[buffer(0)]], - uint tid [[thread_position_in_grid]]) { - out[tid] = Op::init; + uint2 index [[thread_position_in_grid]], + uint2 grid_dim [[threads_per_grid]]) { + out[index.x + grid_dim.x * int64_t(index.y)] = Op::init; } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/scan.cpp b/Source/Cmlx/mlx-generated/scan.cpp index 1fd1d74d..9a291a0b 100644 --- a/Source/Cmlx/mlx-generated/scan.cpp +++ b/Source/Cmlx/mlx-generated/scan.cpp @@ -273,7 +273,7 @@ struct Power { return {0.0, 0.0}; } auto x_theta = metal::atan2(x.imag, x.real); - auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag); + auto x_ln_r = metal::log(hypot(x.real, x.imag)); auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta); auto phase = y.imag * x_ln_r + y.real * x_theta; return {mag * metal::cos(phase), mag * metal::sin(phase)}; diff --git a/Source/Cmlx/mlx-generated/scatter.cpp b/Source/Cmlx/mlx-generated/scatter.cpp index b8bf7c10..ed991a65 100644 --- a/Source/Cmlx/mlx-generated/scatter.cpp +++ b/Source/Cmlx/mlx-generated/scatter.cpp @@ -77,8 +77,8 @@ METAL_FUNC void scatter_impl( for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) { LocT out_idx = out_offset; for (int i = 0; i < NIDX; ++i) { - auto idx_loc = indices.row_contiguous[i] - ? ind_idx + LocT idx_loc = indices.row_contiguous[i] + ? LocT(ind_idx) : elem_to_loc( ind_idx, &indices.shapes[indices.ndim * i], @@ -89,7 +89,7 @@ METAL_FUNC void scatter_impl( out_idx += static_cast(idx_val) * static_cast(out_strides[ax]); } - auto upd_idx = ind_idx * static_cast(upd_size) + gid.x; + LocT upd_idx = ind_idx * static_cast(upd_size) + gid.x; if constexpr (!UPD_ROW_CONTIG) { upd_idx = elem_to_loc(upd_idx, upd_shape, upd_strides, upd_ndim); } diff --git a/Source/Cmlx/mlx-generated/steel_attention.cpp b/Source/Cmlx/mlx-generated/steel_attention.cpp index 75c2edbd..31469e83 100644 --- a/Source/Cmlx/mlx-generated/steel_attention.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention.cpp @@ -1780,7 +1780,7 @@ struct GEMMKernel { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h" -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. using namespace mlx::steel; @@ -1795,6 +1795,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; struct MaxOp { template @@ -1844,6 +1845,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -1857,6 +1859,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -1883,6 +1886,12 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -1895,13 +1904,15 @@ template < constexpr short LDQ_tgp = BD + padQ; constexpr short LDK_tgp = BK + padK; - constexpr short LDV_tgp = BD + padV; + constexpr short LDV_tgp = BV + padV; constexpr short tgp_mem_0 = (BK + padK) * (BD); - constexpr short tgp_mem_1 = BK * (BD + padV); + constexpr short tgp_mem_1 = BK * (BV + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = (WN == 2) && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -1931,7 +1942,7 @@ template < using VBlockLoader = BlockLoaderT< /* typename T = */ T, /* short BROWS = */ BK, - /* short BCOLS = */ BD, + /* short BCOLS = */ BV, /* short kDstStrRow = */ LDV_tgp, /* short kDstStrCol = */ 1, /* short reduction_dim = */ 0, @@ -1949,8 +1960,15 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; + + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert( + BD % WN == 0 && BV % WN == 0, + "The head dims must split evenly across WN"); - constexpr int kNWarps = WM * WN; + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -1961,14 +1979,19 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + constexpr int TVh = BV / (kFragSize * WN); + constexpr int BVh = BV / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -1976,22 +1999,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BVh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -2043,17 +2089,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -2154,7 +2234,7 @@ template < // Load V blocks if (!align_K && kb == (params->NK_aligned)) { - loader_v.load_safe(short2(BD, params->kL_rem)); + loader_v.load_safe(short2(BV, params->kL_rem)); } else { loader_v.load_unsafe(); } @@ -2206,7 +2286,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TVh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -2241,11 +2321,25 @@ template < Otile.template row_bin_op(sum_score); threadgroup_barrier(mem_flags::mem_none); + // Store the logsumexp for the backward pass + if (save_lse && sn == 0) { + using stile_t = decltype(Stile); + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + sm + i * stile_t::kFragRows; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = M_LN2_F * (max_score[i] + metal::log2(sum_score[i])); + } + } + } + // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BVh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BVh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; diff --git a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp index df083170..222bb185 100644 --- a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp @@ -1290,7 +1290,7 @@ struct BlockSwizzle { /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h" -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. using namespace mlx::steel; @@ -1305,6 +1305,7 @@ constant bool align_K [[function_constant(201)]]; constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; +constant bool save_lse [[function_constant(303)]]; template struct TransformScale { @@ -1364,6 +1365,7 @@ template < int BQ, int BK, int BD, + int BV, int WM, int WN, typename MaskType = float, @@ -1377,6 +1379,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* lse [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -1404,6 +1407,13 @@ template < tidl.y * params->O_strides[1] + // Head tidl.x * BQ * params->O_strides[2]; // Sequence + if (save_lse) { + // [B, H, qL], contiguous. + lse += tidl.z * params->H * params->qL + // Batch + tidl.y * params->qL + // Head + tidl.x * BQ; // Sequence + } + if (has_mask) { mask += tidl.z * mask_params->M_strides[0] + // Batch tidl.y * mask_params->M_strides[1]; // Head @@ -1424,11 +1434,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -1709,8 +1721,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } @@ -1760,6 +1772,16 @@ template < Otile.template row_bin_op(rcp); + if (save_lse && sn == 0) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const short r = tm + i * otile_t::kFragRowsJump + sm; + if (align_Q || !is_last_q || r < params->qL_rem) { + lse[r] = max_score[i] * M_LN2_F + metal::log(sum_score[i]); + } + } + } + // Store results O += tm * int(params->O_strides[2]); @@ -1777,15 +1799,10 @@ template < // Head-dim split attention kernel /////////////////////////////////////////////////////////////////////////////// -// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup -// accumulator working set of attention_nax (TD output fragments plus the S -// fragments) is what gates tensor-unit throughput, so this kernel splits the -// head dim across the WN = 2 simdgroups of the second warp dimension: each -// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, -// halving its accumulator set. The pair exchanges its partial Q@K.T sums -// through threadgroup memory, then both simdgroups run softmax redundantly on -// the full S tile (the row statistics are cheap) and each accumulates P@V for -// its own half of Dv. +// Split wide heads across WN simdgroups to reduce the accumulator working set. +// Each group owns one D slice. The groups exchange their partial Q@K.T scores +// through threadgroup memory and reduce them in the same order. Each group then +// runs softmax and accumulates its P@V slice. // clang-format off template < @@ -1806,6 +1823,7 @@ template < const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], const device MaskType* mask [[buffer(6), function_constant(has_mask)]], const device T* sinks [[buffer(7), function_constant(has_sinks)]], + device float* LSE [[buffer(8), function_constant(save_lse)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]], @@ -1844,9 +1862,9 @@ template < constexpr short kU = 16; // The WM simdgroups along the first warp dimension split the Q sequence; - // the WN simdgroups along the second split the head dim. The exchange - // below reduces exactly one peer, so WN is fixed at 2. - static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + // the WN simdgroups along the second split the head dim. + static_assert( + WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4"); constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, @@ -1867,20 +1885,20 @@ template < constexpr int BDh = BD / WN; static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); - static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + static_assert(TK % 2 == 0, "S fragments are computed pair by pair"); const short row_group = simd_group_id / WN; - const short d_half = simd_group_id % WN; + const short d_group = simd_group_id % WN; using otile_t = NAXTile; otile_t Otile; Otile.clear(); const short tm = kU * TQ * row_group; - Q += tm * int(params->Q_strides[2]) + d_half * BDh; - K += d_half * BDh; - V += d_half * BDh; - O += tm * int(params->O_strides[2]) + d_half * BDh; + Q += tm * int(params->Q_strides[2]) + d_group * BDh; + K += d_group * BDh; + V += d_group * BDh; + O += tm * int(params->O_strides[2]) + d_group * BDh; constexpr short kRowsPT = otile_t::kRowsPerThread; @@ -1920,12 +1938,11 @@ template < using stile_t = NAXTile; constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; - // One slot per (row group, half): a fragment pair in per-lane-linear - // layout. Both halves share the fragment-to-lane mapping, so the - // exchange needs no coordinate math. - threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + // One slot per (row group, D group) in per-lane-linear layout. All groups + // share the fragment-to-lane mapping. + threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32]; - // Keep the simdgroup's Q half resident in registers for the whole KV + // Keep the simdgroup's Q slice resident in registers for the whole KV // loop: TDh fragments of T are cheap next to the accumulators. NAXTile Qtiles[TDh]; STEEL_PRAGMA_UNROLL @@ -1950,7 +1967,7 @@ template < stile_t Stile; Stile.clear(); - // S = Q @ K.T, this half of D only, exchanged pair by pair. + // S = Q @ K.T for this D slice. STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { STEEL_PRAGMA_UNROLL @@ -1974,26 +1991,37 @@ template < Ktile.frag_at(1, 0), metal::true_type{}); } + } - // Exchange the partial pair and reduce. - threadgroup AccumType* slot = s_xchg[row_group][d_half]; - thread auto& s0 = Stile.frag_at(0, ik); - thread auto& s1 = Stile.frag_at(0, ik + 1); - const short base = short(simd_lane_id) * (2 * kEPF); + // Exchange all partial scores and reduce them in a fixed order. + threadgroup AccumType* slot = s_xchg[row_group][d_group]; + const short base = short(simd_lane_id) * (TK * kEPF); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - slot[base + i] = s0[i]; - slot[base + kEPF + i] = s1[i]; + slot[base + ik * kEPF + i] = s[i]; } - threadgroup_barrier(mem_flags::mem_threadgroup); - const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& s = Stile.frag_at(0, ik); STEEL_PRAGMA_UNROLL for (short i = 0; i < kEPF; i++) { - s0[i] += peer[base + i]; - s1[i] += peer[base + kEPF + i]; + s[i] = s_xchg[row_group][0][base + ik * kEPF + i]; + } + STEEL_PRAGMA_UNROLL + for (short peer_group = 1; peer_group < WN; peer_group++) { + const threadgroup AccumType* peer = s_xchg[row_group][peer_group]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s[i] += peer[base + ik * kEPF + i]; + } } - threadgroup_barrier(mem_flags::mem_threadgroup); } + threadgroup_barrier(mem_flags::mem_threadgroup); // Scale S STEEL_PRAGMA_UNROLL @@ -2114,7 +2142,7 @@ template < } } - // Do softmax (redundantly per half; the row statistics are cheap) + // Do softmax in each D group; the row statistics are cheap. metal::vec new_max; metal::vec factor; STEEL_PRAGMA_UNROLL @@ -2142,7 +2170,7 @@ template < simdgroup_barrier(mem_flags::mem_none); - // O = P @ V, this half of Dv only. + // O = P @ V for this D slice. STEEL_PRAGMA_UNROLL for (short id = 0; id < TDh; id += 2) { STEEL_PRAGMA_UNROLL @@ -2183,6 +2211,21 @@ template < rcp[i] = 1.f / sum_score[i]; } + if (save_lse && d_group == 0 && sn == 0) { + const int lse_row_base = int(tid.x) * BQ + tm; + const int lse_head_off = (int(tid.z) * params->H + int(tid.y)) * params->qL; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + const int row = lse_row_base + i * stile_t::kFragRowsJump + sm; + if (row < params->qL) { + LSE[lse_head_off + row] = (sum_score[i] == 0) + ? -INFINITY + : float(M_LN2_F * (max_score[i] + metal::log2(sum_score[i]))); + } + } + } + Otile.template row_bin_op(rcp); if (!align_Q && is_last_q) { diff --git a/Source/Cmlx/mlx-generated/steel_conv.cpp b/Source/Cmlx/mlx-generated/steel_conv.cpp index 2af0e560..85aca402 100644 --- a/Source/Cmlx/mlx-generated/steel_conv.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv.cpp @@ -59,7 +59,7 @@ implicit_gemm_conv_2d( using loader_a_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DInputBlockLoaderSmallChannels< @@ -97,7 +97,7 @@ implicit_gemm_conv_2d( // Weight loader using loader_b_t = typename metal::conditional_t< // Check for small channel specialization - N_CHANNELS != 0 && N_CHANNELS <= 4, + N_CHANNELS != 0 && N_CHANNELS <= 8, // Go to small channel specialization Conv2DWeightBlockLoaderSmallChannels< diff --git a/Source/Cmlx/mlx-generated/steel_conv_general.cpp b/Source/Cmlx/mlx-generated/steel_conv_general.cpp index 0024a77f..a3e98ac4 100644 --- a/Source/Cmlx/mlx-generated/steel_conv_general.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv_general.cpp @@ -71,8 +71,8 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const device T* src[n_rows]; @@ -271,8 +271,8 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh; const short base_ww; - short weight_h; - short weight_w; + int weight_h; + int weight_w; const int start_row; @@ -595,7 +595,7 @@ implicit_gemm_conv_2d_general( if (offset_n >= gemm_params->N) return; - short diff = gemm_params->N - offset_n; + int diff = gemm_params->N - offset_n; STEEL_PRAGMA_UNROLL for (int i = 0; i < mma_t::TM; i++) { diff --git a/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp b/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp index 3219ac62..bd6979f6 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_gather.cpp @@ -33,9 +33,10 @@ template < [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { @@ -57,8 +58,19 @@ template < using loader_b_t = typename gemm_kernel::loader_b_t; using mma_t = typename gemm_kernel::mma_t; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } @@ -66,188 +78,127 @@ template < threadgroup T As[gemm_kernel::tgp_mem_size_a]; threadgroup T Bs[gemm_kernel::tgp_mem_size_b]; - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; const size_t c_row_long = size_t(c_row); const size_t c_col_long = size_t(c_col); // Prepare threadgroup bounds - const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row)); const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col)); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[c_row]; - short offset_next = 0; - int n = 0; - while (n < tgp_bm) { - n++; - offset = offset_next; - index = index_next; - offset_next = tgp_bm; - for (; n < tgp_bm; n++) { - if (rhs_indices[c_row + n] != index) { - offset_next = n; - index_next = rhs_indices[c_row + n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); + // Prepare threadgroup mma operation + thread mma_t mma_op(simd_group_id, simd_lane_id); - // Prepare threadgroup mma operation - thread mma_t mma_op(simd_group_id, simd_lane_id); + // Prepare threadgroup loading operations + thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); + thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id); + + // Prepare iterations + const int gemm_k_iterations = params->gemm_k_iterations_aligned; - // Prepare threadgroup loading operations - thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id); - thread loader_b_t loader_b( - B + index * params->batch_stride_b, - params->ldb, + // Do unaligned K iterations first + if (!align_K) { + const int k_last = params->gemm_k_iterations_aligned * BK; + const int k_remain = params->K - k_last; + const size_t k_jump_a = + transpose_a ? params->lda * size_t(k_last) : size_t(k_last); + const size_t k_jump_b = + transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); + + // Move loader source ahead to end + loader_a.src += k_jump_a; + loader_b.src += k_jump_b; + + // Load tile + const short2 tile_dims_A = + transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); + const short2 tile_dims_B = + transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); + + loader_a.load_safe(tile_dims_A); + loader_b.load_safe(tile_dims_B); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Do matmul + mma_op.mma(As, Bs); + + // Reset source back to start + loader_a.src -= k_jump_a; + loader_b.src -= k_jump_b; + } + + const short lbk = 0; + + // Tile aligned don't check + if ((tgp_bm == BM) && (align_N || tgp_bn == BN)) { + gemm_kernel::gemm_loop( + As, Bs, - simd_group_id, - simd_lane_id); - - // Prepare iterations - const int gemm_k_iterations = params->gemm_k_iterations_aligned; - - // Do unaligned K iterations first - if (!align_K) { - const int k_last = params->gemm_k_iterations_aligned * BK; - const int k_remain = params->K - k_last; - const size_t k_jump_a = - transpose_a ? params->lda * size_t(k_last) : size_t(k_last); - const size_t k_jump_b = - transpose_b ? size_t(k_last) : params->ldb * size_t(k_last); - - // Move loader source ahead to end - loader_a.src += k_jump_a; - loader_b.src += k_jump_b; - - // Load tile - const short2 tile_dims_A = - transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm); - const short2 tile_dims_B = - transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain); - - loader_a.load_safe(tile_dims_A); - loader_b.load_safe(tile_dims_B); + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - threadgroup_barrier(mem_flags::mem_threadgroup); + // Tile partially aligned check rows + else if (align_N || tgp_bn == BN) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Do matmul - mma_op.mma(As, Bs); + // Tile partially aligned check cols + else if (tgp_bm == BM) { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Reset source back to start - loader_a.src -= k_jump_a; - loader_b.src -= k_jump_b; - } + // Nothing aligned so check both rows and cols + else { + gemm_kernel::gemm_loop( + As, + Bs, + gemm_k_iterations, + loader_a, + loader_b, + mma_op, + tgp_bm, + tgp_bn, + lbk, + LoopAlignment{}); + } - // Matrix level aligned never check - if (align_M && align_N) { - for (int k = 0; k < gemm_k_iterations; k++) { - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Load elements into threadgroup - loader_a.load_unsafe(); - loader_b.load_unsafe(); - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // Multiply and accumulate threadgroup elements - mma_op.mma(As, Bs); - - // Prepare for next iteration - loader_a.next(); - loader_b.next(); - } - - // Store results to device memory - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } else { - const short lbk = 0; - - // Tile aligned don't check - if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - if (offset_next - offset == BM) { - mma_op.store_result(C, params->ldd); - } else { - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - } - - // Tile partially aligned check rows - else if (align_N || tgp_bn == BN) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(BN, offset_next)); - } - - // Tile partially aligned check cols - else if (align_M || tgp_bm == BM) { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - - // Nothing aligned so check both rows and cols - else { - gemm_kernel::gemm_loop( - As, - Bs, - gemm_k_iterations, - loader_a, - loader_b, - mma_op, - tgp_bm, - tgp_bn, - lbk, - LoopAlignment{}); - mma_op.store_result_slice( - C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next)); - } - } + // Store results to device memory + if (tgp_bm == BM && (align_N || tgp_bn == BN)) { + mma_op.store_result(C, params->ldd); + } else { + mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm)); } } diff --git a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp index 7312ba45..8ff7cdb9 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp @@ -33,9 +33,11 @@ template < gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], - const device uint32_t* rhs_indices [[buffer(2)]], + const device int32_t* offsets [[buffer(2)]], device T* C [[buffer(3)]], const constant GEMMParams* params [[buffer(4)]], + const constant int& num_groups [[buffer(5)]], + uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { constexpr short SM = BM / WM; @@ -44,108 +46,68 @@ gather_mm_rhs_nax( constexpr short TM = SM / 16; constexpr short TN = SN / 16; + int c_row; + int group; + short tgp_bm; if (params->tiles_n <= static_cast(tid.x) || - params->tiles_m <= static_cast(tid.y)) { + !schedule_row_tile( + offsets, + num_groups, + params->M, + tid.y, + simd_lane_id, + c_row, + group, + tgp_bm)) { return; } - // Find the block in A, B, C - const int c_row = tid.y * BM; const int c_col = tid.x * BN; - const size_t c_row_long = size_t(c_row); - const size_t c_col_long = size_t(c_col); + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + const short sgp_sm = short(clamp(tgp_bm - tm, 0, int(SM))); + const short sgp_sn = + align_N ? SN : short(clamp(params->N - (c_col + tn), 0, int(SN))); + const size_t c_row_long = size_t(c_row + tm); + const size_t c_col_long = size_t(c_col + tn); A += transpose_a ? c_row_long : c_row_long * params->lda; + B += group * params->batch_stride_b; B += transpose_b ? c_col_long * params->ldb : c_col_long; C += c_row_long * params->ldd + c_col_long; - rhs_indices += c_row; - - const short tm = SM * (simd_group_id / WN); - const short tn = SN * (simd_group_id % WN); - - const int sgp_sm_int = - align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); - const short sgp_sm = short(sgp_sm_int); - const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); - const int sgp_sn_int = - align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); - const short sgp_sn = short(sgp_sn_int); - const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); - - A += transpose_a ? tm : (tm * params->lda); - B += transpose_b ? (tn * params->ldb) : tn; - C += tm * params->ldd + tn; - rhs_indices += tm; - - // Do as many matmuls as necessary - uint32_t index; - short offset; - uint32_t index_next = rhs_indices[0]; - short offset_next = 0; - int n = 0; - while (n < sgp_sm) { - n++; - offset = offset_next; - index = index_next; - offset_next = sgp_sm; - for (; n < sgp_sm; n++) { - if (rhs_indices[n] != index) { - offset_next = n; - index_next = rhs_indices[n]; - break; - } - } - threadgroup_barrier(mem_flags::mem_none); - - NAXTile Ctile; - - dispatch_bool(align_K, [&](auto kAlignedK) { - dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { - dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K - T, - SM, - SN, - SK, - BK, - transpose_a, - transpose_b, - kAlignedM.value, - kAlignedN.value, - kAlignedK.value, - AccumType>; - Ctile = do_gemm( - A, - B + index * params->batch_stride_b, - params->lda, - params->ldb, - params->K, - params->gemm_k_iterations_aligned, - sgp_sm, - sgp_sn); - - if constexpr (kAlignedN.value) { - if (offset_next - offset == SM) { - Ctile.store(C, int(params->ldd)); - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(SN, offset_next)); - } - } else { - Ctile.store_slice( - C, - int(params->ldd), - short2(0, offset), - short2(sgp_sn, offset_next)); - } - }); + dispatch_bool(align_K, [&](auto kAlignedK) { + dispatch_bool(sgp_sm == SM, [&](auto kAlignedM) { + dispatch_bool(sgp_sn == SN, [&](auto kAlignedN) { + NAXTile Ctile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + params->K, + params->gemm_k_iterations_aligned, + sgp_sm, + sgp_sn); + + if constexpr (kAlignedM.value && kAlignedN.value) { + Ctile.store(C, int(params->ldd)); + } else { + Ctile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } }); }); - } + }); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp b/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp index 84842823..1e1d3726 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_segmented.cpp @@ -92,8 +92,8 @@ template < k_start = segments[tid.z]; k_end = segments[tid.z + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid.z * params->batch_stride_d; // Prepare threadgroup mma operation diff --git a/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp index 519087e1..2e7adfb0 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp @@ -64,8 +64,8 @@ void segmented_mm_nax( k_start = segments[tid_s]; k_end = segments[tid_s + 1]; } - A += transpose_a ? k_start * params->lda : k_start; - B += transpose_b ? k_start : k_start * params->ldb; + A += transpose_a ? int64_t(k_start) * params->lda : k_start; + B += transpose_b ? k_start : int64_t(k_start) * params->ldb; C += tid_s * params->batch_stride_d; constexpr short SM = BM / WM; diff --git a/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp new file mode 100644 index 00000000..d9253fc8 --- /dev/null +++ b/Source/Cmlx/mlx-generated/steel_gemm_thin_nax.cpp @@ -0,0 +1,242 @@ +namespace mlx::core::metal { + +const char* steel_gemm_thin_nax() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h" +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool use_out_source [[function_constant(100)]]; +constant bool do_axpby [[function_constant(110)]]; + +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; +constant bool align_K [[function_constant(202)]]; + +// Each simdgroup computes one (16 * UM) x (16 * UN) block of D. It keeps a K +// chunk of A in registers and MPP reads the B chunk from device memory. +// KS groups of SM x SN simdgroups split K. +// clang-format off +template < + typename T, + int UM, + int UN, + int SM, + int SN, + int KS, + bool transpose_b> +[[kernel, max_total_threads_per_threadgroup(SM * SN * KS * 32)]] void gemm_thin_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device T* C [[buffer(2), function_constant(use_out_source)]], + device T* D [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on + using extents_t = metal::dextents; + using strides_t = metal::array; + using tensor_t = metal::tensor; + + constexpr int TSM = 16 * UM; + constexpr int TSN = 16 * UN; + constexpr int BK = 64; + + const int M = params->M; + const int N = params->N; + const int K = params->K; + const int lda = params->lda; + const int ldb = params->ldb; + + // Threadgroups go along M first, so groups that run together share B. + const int tgp = tid.x; + const int tile_m = tgp % params->tiles_m; + const int tile_n = tgp / params->tiles_m; + + const int k_group = simd_group_id / (SM * SN); + const int sg = simd_group_id % (SM * SN); + int m_off = (tile_m * SM + sg % SM) * TSM; + int n_off = (tile_n * SN + sg / SM) * TSN; + + // A block at the edge moves back inside if the matrix is large enough. It + // then writes again some outputs of a neighbor block, with equal values. + const bool active = (align_M || m_off < M) && (align_N || n_off < N); + if (!align_M && m_off + TSM > M && M >= TSM) { + m_off = M - TSM; + } + if (!align_N && n_off + TSN > N && N >= TSN) { + n_off = N - TSN; + } + const bool inside = + (align_M || m_off + TSM <= M) && (align_N || n_off + TSN <= N); + const int rows = active ? M - m_off : 0; + const int cols = active ? N - n_off : 0; + + A += int64_t(m_off) * lda; + B += transpose_b ? int64_t(n_off) * ldb : int64_t(n_off); + D += int64_t(m_off) * params->ldd + n_off; + if (use_out_source) { + C += + int64_t(m_off) * addmm_params->ldc + int64_t(n_off) * addmm_params->fdc; + } + + tensor_t tA((device T*)A, extents_t(K, rows), strides_t{1, lda}); + tensor_t tB( + (device T*)B, + transpose_b ? extents_t(K, cols) : extents_t(cols, K), + strides_t{1, ldb}); + + auto a_chunk = [&](int k) { return tA.template slice(k, 0); }; + auto b_chunk = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + auto b_tail = [&](int k) { + if constexpr (transpose_b) { + return tB.template slice(k, 0); + } else { + return tB.template slice(0, k); + } + }; + using a_chunk_t = decltype(a_chunk(0)); + using b_chunk_t = decltype(b_chunk(0)); + + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + TSM, + TSN, + BK, + false, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + mpp::tensor_ops::matmul2d op; + + // The accumulator layout depends only on the element types, so the masked + // runs below can also use it. + auto Dacc = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] = 0.0f; + } + auto Areg = op.template get_left_input_cooperative_tensor(); + + const int k_iters = params->gemm_k_iterations_aligned; + for (int it = 0; it < k_iters; it++) { + // Keep the simdgroups near each other in K, so that their shared loads + // hit in cache. + if ((it & 3) == 0) { + threadgroup_barrier(mem_flags::mem_none); + } + const int k = (it * KS + k_group) * BK; + if (!active || k >= K) { + continue; + } + if (inside && (align_K || k + BK <= K)) { + auto a = a_chunk(k); + Areg.load(a); + auto b = b_chunk(k); + op.run(Areg, b, Dacc); + } else if (inside) { + auto a = tA.template slice(k, 0); + auto b = b_tail(k); + op.run(a, b, Dacc); + } else { + // The views stop at the matrix edge and MPP masks the rest. + const int kk = min(BK, K - k); + tensor_t vA((device T*)A + k, extents_t(kk, rows), strides_t{1, lda}); + tensor_t vB( + (device T*)B + (transpose_b ? int64_t(k) : int64_t(k) * ldb), + transpose_b ? extents_t(kk, cols) : extents_t(cols, kk), + strides_t{1, ldb}); + op.run(vA, vB, Dacc); + } + } + + if constexpr (KS > 1) { + // K groups 1 to KS - 1 store their partial sums and group 0 adds them. + using tg_tensor_t = + metal::tensor; + constexpr int kBlockSize = TSM * TSN; + static_assert( + (KS - 1) * SM * SN * kBlockSize <= 4096, + "Partial sums need more than 16 KB of threadgroup memory"); + threadgroup float partials[(KS - 1) * SM * SN * kBlockSize]; + + if (active && k_group > 0) { + tg_tensor_t tP( + partials + ((k_group - 1) * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + Dacc.store(tP); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (active && k_group == 0) { + for (int g = 0; g < KS - 1; g++) { + tg_tensor_t tP( + partials + (g * SM * SN + sg) * kBlockSize, + extents_t(TSN, TSM), + strides_t{1, TSN}); + auto part = op.template get_destination_cooperative_tensor< + a_chunk_t, + b_chunk_t, + float>(); + part.load(tP); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + Dacc[i] += part[i]; + } + } + } + } + if (k_group != 0 || !active) { + return; + } + + auto Dout = + op.template get_destination_cooperative_tensor(); + STEEL_PRAGMA_UNROLL + for (uint16_t i = 0; i < Dacc.get_capacity(); i++) { + float d = Dacc[i]; + if (use_out_source) { + auto idx = Dacc.get_multidimensional_index(i); + const int r = idx[1]; + const int c = idx[0]; + if (inside || (r < rows && c < cols)) { + float c_val = static_cast( + C[int64_t(r) * addmm_params->ldc + int64_t(c) * addmm_params->fdc]); + d = do_axpby ? addmm_params->alpha * d + addmm_params->beta * c_val + : d + c_val; + } + } + Dout[i] = static_cast(d); + } + + tensor_t tD(D, extents_t(cols, rows), strides_t{1, params->ldd}); + if (inside) { + auto d = tD.template slice(0, 0); + Dout.store(d); + } else { + Dout.store(tD); + } +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/unary_ops.cpp b/Source/Cmlx/mlx-generated/unary_ops.cpp index 3a680b7f..4ab9ac5f 100644 --- a/Source/Cmlx/mlx-generated/unary_ops.cpp +++ b/Source/Cmlx/mlx-generated/unary_ops.cpp @@ -349,11 +349,11 @@ struct fp8_e4m3 { } operator float16_t() thread { - uint16_t v = (bits & 127) << 7; - half converted = as_type(v); - converted *= 256.0; - auto sign = bits & 128; - return (sign ? -converted : converted); + uint16_t v = bits & 127; + uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15; + uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit; + half converted = as_type(u); + return converted * 256.0; } operator bfloat16_t() thread { @@ -446,7 +446,7 @@ struct Abs { return x; }; complex64_t operator()(complex64_t x) thread { - return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; + return {hypot(x.real, x.imag), 0}; }; }; @@ -719,7 +719,7 @@ struct Round { struct Sigmoid { template T operator()(T x) thread { - auto y = 1 / (1 + metal::exp(metal::abs(x))); + auto y = 1 / (1 + metal::precise::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -736,8 +736,8 @@ struct Sign { if (x == complex64_t(0)) { return x; } - return x / - (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag); + auto r = hypot(x.real, x.imag); + return {x.real / r, x.imag / r}; }; }; @@ -789,7 +789,7 @@ struct Sqrt { auto b_abs = metal::precise::sqrt((r - x.real) / 2.0); auto b = metal::copysign(b_abs, x.imag); return {a, b}; - } + }; }; struct Rsqrt { diff --git a/Source/Cmlx/mlx-generated/utils.cpp b/Source/Cmlx/mlx-generated/utils.cpp index 18d9644d..047bcd78 100644 --- a/Source/Cmlx/mlx-generated/utils.cpp +++ b/Source/Cmlx/mlx-generated/utils.cpp @@ -681,10 +681,20 @@ constexpr complex_t operator*(complex_t a, complex_t b) { template constexpr complex_t operator/(complex_t a, complex_t b) { - auto denom = b.real * b.real + b.imag * b.imag; - auto x = a.real * b.real + a.imag * b.imag; - auto y = a.imag * b.real - a.real * b.imag; - return {x / denom, y / denom}; + float ar = a.real, ai = a.imag; + float br = b.real, bi = b.imag; + if (abs(br) >= abs(bi)) { + auto ratio = bi / br; + auto denom = br + bi * ratio; + return { + static_cast((ar + ai * ratio) / denom), + static_cast((ai - ar * ratio) / denom)}; + } + auto ratio = br / bi; + auto denom = bi + br * ratio; + return { + static_cast((ar * ratio + ai) / denom), + static_cast((ai * ratio - ar) / denom)}; } template < @@ -692,11 +702,7 @@ template < typename U, enable_if_t && is_lane_convertible_v, bool> = true> constexpr complex_t operator/(U a, complex_t b) { - auto scalar = static_cast(a); - auto denom = b.real * b.real + b.imag * b.imag; - auto x = scalar * b.real; - auto y = -scalar * b.imag; - return {x / denom, y / denom}; + return complex_t(a) / b; } template @@ -739,6 +745,7 @@ static MTL_CONST constexpr int REDUCE_N_WRITES = 4; static MTL_CONST constexpr int SOFTMAX_N_READS = 4; static MTL_CONST constexpr int RMS_N_READS = 4; static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; +static MTL_CONST constexpr int CROSS_ENTROPY_N_READS = 4; // Instantiate a templated kernel. // Extra args are used as template parameters: @@ -1133,6 +1140,29 @@ inline complex64_t log1p(complex64_t in) { } } +// https://github.com/pytorch/pytorch/blob/a82aae9d4a7827849ce50f31c4c7ee8f278d05f5/c10/metal/utils.h#L554 +inline float hypot(float x, float y) { + if (metal::isinf(x) || metal::isinf(y)) { + return metal::numeric_limits::infinity(); + } + if (metal::isnan(x) || metal::isnan(y)) { + return metal::numeric_limits::quiet_NaN(); + } + float a = metal::fmax(metal::fabs(x), metal::fabs(y)); + float b = metal::fmin(metal::fabs(x), metal::fabs(y)); + if (a == 0.0f) { + return 0.0f; + } + float r = (b / a) * (b / a); + float sqrt_1_plus_r = metal::precise::sqrt(1.0f + r); + float h1 = metal::sqrt(2.0f) * a; + float h2 = a + a * r / 2.0f; + float h3 = a * sqrt_1_plus_r; + bool is_h1 = (a == b); + bool is_h2 = ((sqrt_1_plus_r == 1.0f) && (r > 0.0f)); + return metal::select(metal::select(h3, h2, is_h2), h1, is_h1); +} + /////////////////////////////////////////////////////////////////////////////// // SIMD shuffle ops /////////////////////////////////////////////////////////////////////////////// @@ -1251,6 +1281,53 @@ inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} + /////////////////////////////////////////////////////////////////////////////// )preamble"; } diff --git a/Source/MLX/Distributed.swift b/Source/MLX/Distributed.swift index 87d8f765..47947636 100644 --- a/Source/MLX/Distributed.swift +++ b/Source/MLX/Distributed.swift @@ -76,12 +76,16 @@ public enum MLXDistributed { /// The rank of this process in the group. public var rank: Int { - Int(mlx_distributed_group_rank(ctx)) + var result: Int32 = 0 + mlx_distributed_group_get_rank(&result, ctx) + return Int(result) } /// The number of processes in the group. public var size: Int { - Int(mlx_distributed_group_size(ctx)) + var result: Int32 = 0 + mlx_distributed_group_get_size(&result, ctx) + return Int(result) } /// Split the group into subgroups based on the provided color. diff --git a/Source/MLX/Documentation.docc/Organization/arithmetic.md b/Source/MLX/Documentation.docc/Organization/arithmetic.md index 2a585129..bc9203bd 100644 --- a/Source/MLX/Documentation.docc/Organization/arithmetic.md +++ b/Source/MLX/Documentation.docc/Organization/arithmetic.md @@ -215,7 +215,7 @@ Note: the `-` and `/` operators are not able to be linked here. - ``blockMaskedMM(_:_:blockSize:maskOut:maskLHS:maskRHS:stream:)`` - ``addMM(_:_:_:alpha:beta:stream:)`` - ``quantizedMM(_:_:scales:biases:transpose:groupSize:bits:mode:stream:)`` -- ``gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:sortedIndices:stream:)`` +- ``gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:globalScale:sortedIndices:stream:)`` - ``quantizedQuantizedMM(_:_:scales:groupSize:bits:mode:globalScaleX:globalScaleW:stream:)`` - ``segmentedMM(_:_:segments:stream:)`` - ``inner(_:_:stream:)`` diff --git a/Source/MLX/MLXFast.swift b/Source/MLX/MLXFast.swift index 47e5c86c..2a53343a 100644 --- a/Source/MLX/MLXFast.swift +++ b/Source/MLX/MLXFast.swift @@ -4,6 +4,42 @@ import Cmlx public enum MLXFast { + /// Chunked gated delta network forward pass. + /// + /// - Parameters: + /// - queries: shape (B, T, Hk, Dk) + /// - keys: shape (B, T, Hk, Dk) + /// - values: shape (B, T, Hv, Dv) + /// - gamma: decay rate in linear space (B, T, Hv) + /// - beta: delta update rates (B, T, Hv) + /// - initialState: initial hidden state (B, Hv, Dv, Dk), zeros if not given + /// - mask: optional boolean mask (B, T); masked steps leave the state unchanged + /// - stream: stream or device to evaluate on + /// - Returns: `output` (B, T, Hv, Dv) and the final `state` (B, Hv, Dv, Dk) + public static func gatedDeltaUpdate( + queries: MLXArray, + keys: MLXArray, + values: MLXArray, + gamma: MLXArray, + beta: MLXArray, + initialState: MLXArray? = nil, + mask: MLXArray? = nil, + stream: StreamOrDevice = .default + ) -> (output: MLXArray, state: MLXArray) { + var result = mlx_vector_array_new() + defer { mlx_vector_array_free(result) } + mlx_fast_gated_delta_update( + &result, + queries.ctx, keys.ctx, values.ctx, + gamma.ctx, beta.ctx, + (initialState ?? .mlxNone).ctx, + (mask ?? .mlxNone).ctx, + stream.ctx) + + let tuple = mlx_vector_array_values(result) + return (output: tuple[0], state: tuple[1]) + } + /// Optimized implementation of `NN.RoPE`. /// /// Used like this: diff --git a/Source/MLX/Memory.swift b/Source/MLX/Memory.swift index 7036633c..ead36edf 100644 --- a/Source/MLX/Memory.swift +++ b/Source/MLX/Memory.swift @@ -358,4 +358,21 @@ public enum Memory { mlx_clear_cache() } } + + /// Get the size of the buffers backing the given arrays in bytes. + /// + /// Each unique buffer is counted once. The full allocator size of each buffer + /// is used, which can exceed the logical size of its arrays. The arrays must be + /// evaluated before calling this function. + /// + /// - Parameter arrays: arrays to count + /// - Returns: size in bytes + public static func bufferSize(_ arrays: some Collection) -> Int { + let vector = new_mlx_vector_array(arrays) + defer { mlx_vector_array_free(vector) } + + var result = 0 + mlx_get_array_buffer_size(&result, vector) + return result + } } diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index a1cba744..98b48b7e 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -1928,6 +1928,7 @@ public func take(_ array: MLXArray, _ indices: MLXArray, stream: StreamOrDevice /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axis:stream:)`` /// - ``transposed(_:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` @@ -1941,6 +1942,19 @@ public func transposed( return MLXArray(result) } +/// Transpose the dimensions of the array. +/// +/// - Parameters: +/// - array: input array +/// - axes: Specifies the source axis for each axis in the new array +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - ``matrixTransposed(_:stream:)`` +/// - ``transposed(_:axis:stream:)`` +/// - ``transposed(_:stream:)`` +/// - ``MLXArray/transposed(axes:stream:)`` public func transposed(_ array: MLXArray, _ axes: Int..., stream: StreamOrDevice = .default) -> MLXArray { @@ -1955,6 +1969,7 @@ public func transposed(_ array: MLXArray, _ axes: Int..., stream: StreamOrDevice /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axes:stream:)`` /// - ``transposed(_:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` @@ -1971,6 +1986,7 @@ public func transposed(_ array: MLXArray, axis: Int, stream: StreamOrDevice = .d /// /// ### See Also /// - +/// - ``matrixTransposed(_:stream:)`` /// - ``transposed(_:axes:stream:)`` /// - ``transposed(_:axis:stream:)`` /// - ``MLXArray/transposed(axes:stream:)`` diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index 5a0da9e1..f653918a 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -1491,7 +1491,7 @@ public func gatherMM( @available( *, deprecated, renamed: - "gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:sortedIndices:stream:)" + "gatherQuantizedMM(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:globalScale:sortedIndices:stream:)" ) public func gatherQuantizedMatmul( _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray?, @@ -1528,6 +1528,8 @@ public func gatherQuantizedMatmul( /// - groupSize: The size of the group in `w` that shares a scale and bias. Default is `64` /// - bits: The number of bits occupied by each element in `w`. Default is `4` /// - mode: The quantization mode. Default is `.affine` +/// - globalScale: The per-input float32 scale used for +/// `nvfp4` quantization of `w`. Only supported on the GPU. /// - sortedIndices: Whether the indices are sorted. Default is `false` /// - stream: Stream or device to evaluate on /// @@ -1535,10 +1537,11 @@ public func gatherQuantizedMatmul( /// - /// - ``quantizedMM(_:_:scales:biases:transpose:groupSize:bits:mode:stream:)`` public func gatherQuantizedMM( - _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray?, + _ x: MLXArray, _ w: MLXArray, scales: MLXArray, biases: MLXArray? = nil, lhsIndices: MLXArray? = nil, rhsIndices: MLXArray? = nil, transpose: Bool = true, groupSize: Int? = nil, bits: Int? = nil, mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil, sortedIndices: Bool = false, stream: StreamOrDevice = .default ) -> MLXArray { @@ -1551,7 +1554,10 @@ public func gatherQuantizedMM( &result, x.ctx, w.ctx, scales.ctx, (biases ?? .mlxNone).ctx, (lhsIndices ?? .mlxNone).ctx, (rhsIndices ?? .mlxNone).ctx, transpose, - gs, bits, mode.rawValue, sortedIndices, + gs, bits, + mode.rawValue, + (globalScale ?? .mlxNone).ctx, + sortedIndices, stream.ctx) return MLXArray(result) @@ -2043,6 +2049,22 @@ public func meshGrid( return mlx_vector_array_values(vec) } +/// Transpose the last two dimensions of the array. +/// +/// ### See Also +/// - ``transposed(_:stream:)`` +/// +/// - Parameters: +/// - array: array to transpose +/// - stream: stream or device to evaluate on +public func matrixTransposed( + _ array: MLXArray, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_matrix_transpose(&result, array.ctx, stream.ctx) + return MLXArray(result) +} + /// Element-wise maximum. /// /// Take the element-wise max of two arrays with @@ -2999,6 +3021,163 @@ public func unstacked( return mlx_vector_array_values(result) } +/// Returns the sorted unique elements of the flattened array. +/// +/// The shape of the output does not depend on the values in `array`, so +/// `array` is not evaluated. Unlike NumPy, the size is never inferred from +/// the values and must be provided. +/// +/// Entries past the last unique element hold `fillValue`. +/// A truncated output keeps the smallest `size` unique elements. +/// +/// This op builds on ``sorted(_:stream:)``, so a `Bool` input needs the CPU +/// stream. +/// +/// ### See Also +/// - ``uniqued(_:size:fillValue:returning:stream:)`` +/// +/// - Parameters: +/// - array: input array +/// - size: The size of the output. If the size is smaller than +/// the number of unique elements of `array`, the output is truncated. +/// If it is larger, the output is padded with `fillValue`. +/// - fillValue: optional value to fill empty positions in the result. If +/// not specified this defaults to the first of the sorted unique elements. +/// - stream: stream or device to evaluate on +/// - Returns: sorted unique elements +public func uniqued( + _ array: MLXArray, + size: Int, + fillValue: MLXArray? = nil, + stream: StreamOrDevice = .default +) -> MLXArray { + var vector = mlx_vector_array_new() + defer { mlx_vector_array_free(vector) } + + mlx_unique( + &vector, + array.ctx, size.int32, + false, false, false, + (fillValue ?? .mlxNone).ctx, + stream.ctx) + + let arrays = mlx_vector_array_values(vector) + return arrays[0] +} + +/// Options for `uniqued()`. +/// +/// See ``uniqued(_:size:fillValue:returning:stream:)``. +public struct UniqueOptions: OptionSet, Sendable { + public let rawValue: Int + public init(rawValue: Int) { self.rawValue = rawValue } + + /// Requests the index in the flattened input of the first + /// occurence of each unique element. + public static let index = UniqueOptions(rawValue: 1 << 0) + + /// Requests the indices of the unique array that rebuild + /// the input to `uniqued()`. + public static let inverse = UniqueOptions(rawValue: 1 << 1) + + /// Requests the number of times each unique element occurs in the + /// input. + public static let counts = UniqueOptions(rawValue: 1 << 2) +} + +/// Return value for `uniqued()`. +/// +/// See ``uniqued(_:size:fillValue:returning:stream:)``. +public struct UniqueResult { + + /// Array of unique values of the given size. + public var values: MLXArray + public var index: MLXArray? + public var inverse: MLXArray? + public var counts: MLXArray? + + public init( + values: MLXArray, + index: MLXArray? = nil, inverse: MLXArray? = nil, counts: MLXArray? = nil + ) { + self.values = values + self.index = index + self.inverse = inverse + self.counts = counts + } +} + +/// Returns the sorted unique elements of the flattened array. +/// +/// The shape of the output does not depend on the values in `array`, so +/// `array` is not evaluated. Unlike NumPy, the size is never inferred from +/// the values and must be provided. +/// +/// Entries past the last unique element hold `fillValue` and their +/// count is `0`, so `sum(result.counts > 0)` gives the number of unique +/// elements as long as the output is not truncated. +/// +/// A truncated output keeps the smallest `size` unique elements, and +/// `result.index`, `result.inverse` and `result.counts` +/// then stop describing all of `array`, +/// and `result.inverse` and `result.counts` stop agreeing with each other: +/// the counts of the dropped elements are gone, while their +/// indices in `result.inverse` are clamped to the last entry. +/// +/// This op builds on ``sorted(_:stream:)``, so a `Bool` input needs the CPU +/// stream. +/// +/// ### See Also +/// - ``uniqued(_:size:fillValue:stream:)`` +/// +/// - Parameters: +/// - array: input array +/// - size: The size of the output. If the size is smaller than +/// the number of unique elements of `array`, the output is truncated. +/// If it is larger, the output is padded with `fillValue`. +/// - fillValue: optional value to fill empty positions in the result. If +/// not specified this defaults to the first of the sorted unique elements. +/// - returning: Additional values to return +/// - stream: stream or device to evaluate on +/// - Returns: sorted unique elements +public func uniqued( + _ array: MLXArray, + size: Int, + fillValue: MLXArray? = nil, + returning: UniqueOptions, + stream: StreamOrDevice = .default +) -> UniqueResult { + var vector = mlx_vector_array_new() + defer { mlx_vector_array_free(vector) } + + mlx_unique( + &vector, + array.ctx, size.int32, + returning.contains(.index), + returning.contains(.inverse), + returning.contains(.counts), + (fillValue ?? .mlxNone).ctx, + stream.ctx) + + let arrays = mlx_vector_array_values(vector) + var result = UniqueResult(values: arrays[0]) + + var readIndex = 1 + if returning.contains(.index) { + result.index = arrays[readIndex] + readIndex += 1 + } + if returning.contains(.inverse) { + result.inverse = arrays[readIndex] + readIndex += 1 + } + if returning.contains(.counts) { + result.counts = arrays[readIndex] + readIndex += 1 + } + return result +} + /// Stop gradients from being computed. /// ///The operation is the identity but it prevents gradients from flowing diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift index 6a0102a4..ac560dad 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 7 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 15 import Foundation import MLX @@ -198,4 +200,272 @@ struct GeneratedFastTests { tolerance: .float32) } } + + @Test("gatedDeltaUpdate/fallback/output") + func test_gatedDeltaUpdate_fallback_output() throws { + // Hv != Hk, shape not supported by the kernel + try withIntegrationState(seed: 53947) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: -0.018037354573607445, + minimum: -0.5729220509529114, + maximum: 0.4151778519153595, + absoluteSum: 30.585939407348633, + positionChecksum: 17.2472412109375, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -0.05421087518334389, -0.12156698107719421, -0.15238893032073975, + -0.14626511931419373, 0.2740327715873718, -0.10090052336454391, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/state") + func test_gatedDeltaUpdate_fallback_state() throws { + // Hv != Hk, shape not supported by the kernel + try withIntegrationState(seed: 93688) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: 0.006113756913691759, + minimum: -0.15103420615196228, + maximum: 0.1642816811800003, + absoluteSum: 18.600929260253906, + positionChecksum: 9.371909141540527, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.07339409738779068, 0.03697947785258293, 0.03930574283003807, + 0.003275439143180847, 0.005969321355223656, 0.0017838235944509506, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/initialState/output") + func test_gatedDeltaUpdate_fallback_initialState_output() throws { + try withIntegrationState(seed: 34247) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: 0.05453367158770561, + minimum: -3.9586031436920166, + maximum: 7.431026935577393, + absoluteSum: 279.72869873046875, + positionChecksum: 135.665087890625, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -1.2578498125076294, 1.2448211908340454, 0.4136521518230438, + -0.41495442390441895, 0.16835975646972656, -0.2699488699436188, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/initialState/state") + func test_gatedDeltaUpdate_fallback_initialState_state() throws { + try withIntegrationState(seed: 95025) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: -0.009089105762541294, + minimum: -0.5524362325668335, + maximum: 0.5001603960990906, + absoluteSum: 72.35923767089844, + positionChecksum: 33.11988830566406, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.20271210372447968, -0.2317522168159485, 0.268993079662323, + 0.03319134935736656, 0.012991426512598991, 0.024516910314559937, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/mask/output") + func test_gatedDeltaUpdate_fallback_mask_output() throws { + // masked steps leave the state unchanged and output zeros + try withIntegrationState(seed: 92033) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MLXArray( + [true, true, true, true, true, true, true, true, false, false], [2, 5]) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState, mask: mask + ).output + expectSummary( + result, + ArraySummary( + shape: [2, 5, 4, 8], + dtype: .float32, + mean: 0.052943795919418335, + minimum: -6.441189765930176, + maximum: 6.602780342102051, + absoluteSum: 320.8973693847656, + positionChecksum: 117.5300537109375, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -5.145998001098633, -0.6544278860092163, 0.9613724946975708, + 0.406800240278244, -0.6966778635978699, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/fallback/mask/state") + func test_gatedDeltaUpdate_fallback_mask_state() throws { + // masked steps leave the state unchanged and output zeros + try withIntegrationState(seed: 40800) { + let queries = MLXRandom.normal([2, 5, 2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [2, 5, 2, 8], dtype: .float32) + let values = MLXRandom.normal([2, 5, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [2, 5, 4], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [2, 5, 4], dtype: .float32) + let initialState = MLXRandom.normal( + [2, 4, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MLXArray( + [true, true, true, true, true, true, true, true, false, false], [2, 5]) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState, mask: mask + ).state + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: -0.005854031536728144, + minimum: -1.9470419883728027, + maximum: 2.252192974090576, + absoluteSum: 151.9166259765625, + positionChecksum: 87.9920883178711, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.21978959441184998, -0.0349111407995224, 0.3030872046947479, + 0.8860476016998291, -0.3838585615158081, 0.10365576297044754, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/kernel/output") + func test_gatedDeltaUpdate_kernel_output() throws { + // a shape the metal kernel supports + try withIntegrationState(seed: 91400) { + let queries = MLXRandom.normal([1, 3, 16, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [1, 3, 16, 128], dtype: .float32) + let values = MLXRandom.normal([1, 3, 32, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [1, 3, 32], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [1, 3, 32], dtype: .float32) + let initialState = MLXRandom.normal( + [1, 32, 128, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).output + expectSummary( + result, + ArraySummary( + shape: [1, 3, 32, 128], + dtype: .float32, + mean: -0.07715722918510437, + minimum: -33.80734634399414, + maximum: 35.589317321777344, + absoluteSum: 67446.203125, + positionChecksum: 29519.645833333332, + sampleIndices: [0, 2457, 4915, 7372, 9830, 12287], + samples: [ + 6.1729230880737305, -1.2932864427566528, 4.83355188369751, + 9.386662483215332, 3.525575637817383, -4.2895827293396, + ]), + tolerance: .float32) + } + } + + @Test("gatedDeltaUpdate/kernel/state") + func test_gatedDeltaUpdate_kernel_state() throws { + // a shape the metal kernel supports + try withIntegrationState(seed: 81342) { + let queries = MLXRandom.normal([1, 3, 16, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.uniform(low: -0.1, high: 0.1, [1, 3, 16, 128], dtype: .float32) + let values = MLXRandom.normal([1, 3, 32, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let gamma = MLXRandom.uniform(low: 0.5, high: 1.0, [1, 3, 32], dtype: .float32) + let beta = MLXRandom.uniform(low: 0.1, high: 0.9, [1, 3, 32], dtype: .float32) + let initialState = MLXRandom.normal( + [1, 32, 128, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXFast.gatedDeltaUpdate( + queries: queries, keys: keys, values: values, gamma: gamma, beta: beta, + initialState: initialState + ).state + expectSummary( + result, + ArraySummary( + shape: [1, 32, 128, 128], + dtype: .float32, + mean: -0.0014557009562849998, + minimum: -3.1630849838256836, + maximum: 2.9891040325164795, + absoluteSum: 172419.046875, + positionChecksum: 84804.59375, + sampleIndices: [0, 104857, 209715, 314572, 419430, 524287], + samples: [ + 0.1790539026260376, -0.17156574130058289, -0.12500464916229248, + -0.21596695482730865, 0.10746431350708008, -0.5037912130355835, + ]), + tolerance: .float32) + } + } } diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift index 11ce7fff..40a3feb4 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 22 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 25 import Foundation import MLX @@ -568,4 +570,112 @@ struct GeneratedQuantizationTests { tolerance: .float32) } } + + @Test("gatherQuantizedMM/nvfp4/globalScale") + func test_gatherQuantizedMM_nvfp4_globalScale() throws { + // global scale is gpu only and needs indices + try withIntegrationState(seed: 48073) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 64], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([2, 0, 3] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, globalScale: globalScale) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 32], + dtype: .float32, + mean: -0.9819647669792175, + minimum: -165.96966552734375, + maximum: 144.31866455078125, + absoluteSum: 5297.81396484375, + positionChecksum: 3304.0732421875, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + -55.56662368774414, -11.178827285766602, -5.531567096710205, + 0.3299040198326111, -16.427204132080078, 57.386314392089844, + ]), + tolerance: .float32) + } + } + + @Test("gatherQuantizedMM/nvfp4/globalScale/transpose") + func test_gatherQuantizedMM_nvfp4_globalScale_transpose() throws { + try withIntegrationState(seed: 33918) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 32], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([2, 0, 3] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, transpose: false, mode: .nvfp4, + globalScale: globalScale) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 64], + dtype: .float32, + mean: -1.6645965576171875, + minimum: -119.096923828125, + maximum: 92.36067199707031, + absoluteSum: 7135.9306640625, + positionChecksum: 4085.1412760416665, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -43.466453552246094, 8.144684791564941, -3.2288098335266113, + 6.519178867340088, 6.189074516296387, 76.47562408447266, + ]), + tolerance: .float32) + } + } + + @Test("gatherQuantizedMM/nvfp4/globalScale/sorted") + func test_gatherQuantizedMM_nvfp4_globalScale_sorted() throws { + try withIntegrationState(seed: 24933) { + let w = + MLXRandom.normal([4, 32, 64]) + * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1) + let globalScale = MLX.abs(w).max(axes: [1, 2]) + let q = (0 ..< 4).map { + MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) + } + let wq = MLX.stacked(q.map { $0.wq }) + let scales = MLX.stacked(q.map { $0.scales }) + let x = MLXRandom.normal([3, 2, 64], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXArray([0, 2, 2] as [UInt32]) + let result = MLX.gatherQuantizedMM( + x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, globalScale: globalScale, + sortedIndices: true) + expectSummary( + result, + ArraySummary( + shape: [3, 2, 32], + dtype: .float32, + mean: -1.1882468461990356, + minimum: -89.71018981933594, + maximum: 76.8901138305664, + absoluteSum: 3490.58984375, + positionChecksum: 2165.3154296875, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + 2.895949125289917, -2.178328037261963, 52.326904296875, 13.967162132263184, + -27.429550170898438, -89.71018981933594, + ]), + tolerance: .float32) + } + } } diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift index bed68264..f5ca87e1 100644 --- a/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift +++ b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift @@ -7,10 +7,12 @@ // // python3 tools/integration_tests/generate.py // -// python mlx: 0.32.2.dev20260910+1f8e74e3f -// vendored mlx (Cmlx): 0.32.2 +// python mlx: 0.32.4.dev20261021+36969c094 +// vendored mlx (Cmlx): 0.32.4 // generator revision: 2 -// cases: 49 +// MLX_ENABLE_TF32: 0 +// device: Apple M5 Max (applegpu_g17s) +// cases: 69 import Foundation import MLX @@ -1169,4 +1171,440 @@ struct GeneratedShapeTests { tolerance: .float32) } } + + @Test("matrixTransposed") + func test_matrixTransposed() throws { + try withIntegrationState(seed: 94372) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matrixTransposed(a) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.33135515451431274, + minimum: -1.0816655158996582, + maximum: 1.7783619165420532, + absoluteSum: 8.355605125427246, + positionChecksum: 4.310164133707683, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.028264261782169342, 0.2258906066417694, 1.7783619165420532, + 0.21972355246543884, 1.0169203281402588, -0.8231726884841919, + ]), + tolerance: .float32) + } + } + + @Test("matrixTransposed/4d") + func test_matrixTransposed_4d() throws { + // only the last two axes move, unlike transposed(a) which reverses all of them + try withIntegrationState(seed: 89528) { + let a = MLXRandom.normal([2, 3, 4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matrixTransposed(a) + expectSummary( + result, + ArraySummary( + shape: [2, 3, 5, 4], + dtype: .float32, + mean: -0.06181011348962784, + minimum: -2.80637788772583, + maximum: 2.895658016204834, + absoluteSum: 101.28848266601562, + positionChecksum: 53.0994384765625, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + -1.9200016260147095, 0.419328898191452, -0.08306543529033661, + -1.4899338483810425, -1.7050701379776, 0.0025647925212979317, + ]), + tolerance: .float32) + } + } + + @Test("uniqued") + func test_uniqued() throws { + try withIntegrationState(seed: 57697) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 6) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .int32, + mean: 2.5, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 15.0, + positionChecksum: 11.666666666666666, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), + tolerance: .exact) + } + } + + @Test("uniqued/padded") + func test_uniqued_padded() throws { + // no fillValue: pins down the default padding + try withIntegrationState(seed: 24389) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.5, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 15.0, + positionChecksum: 7.0, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/fillValue") + func test_uniqued_fillValue() throws { + try withIntegrationState(seed: 96023) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, fillValue: MLXArray(-1)) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.100000023841858, + minimum: -1.0, + maximum: 5.0, + absoluteSum: 19.0, + positionChecksum: 10.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, -1.0, -1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/truncated") + func test_uniqued_truncated() throws { + // keeps the smallest `size` unique elements + try withIntegrationState(seed: 55506) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3) + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .int32, + mean: 1.0, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 3.0, + positionChecksum: 2.6666666666666665, + sampleIndices: [0, 1, 2], + samples: [0.0, 1.0, 2.0]), + tolerance: .exact) + } + } + + @Test("uniqued/float") + func test_uniqued_float() throws { + try withIntegrationState(seed: 21512) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.uniqued(a, size: 12) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: -0.4238990545272827, + minimum: -2.4105277061462402, + maximum: 1.042361855506897, + absoluteSum: 8.695816040039062, + positionChecksum: 3.2035303115844727, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.4105277061462402, -0.8730838298797607, -0.5980460047721863, + -0.14297938346862793, 0.14883598685264587, 1.042361855506897, + ]), + tolerance: .float32) + } + } + + @Test("uniqued/bool") + func test_uniqued_bool() throws { + // bool needs the cpu stream (builds on sort) + try withIntegrationState(seed: 40846) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.uniqued(a, size: 2, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [2], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0, 1], + samples: [0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/values") + func test_uniqued_returning_values() throws { + try withIntegrationState(seed: 21630) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued( + a, size: 10, fillValue: MLXArray(-1), returning: [.index, .inverse, .counts] + ).values + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 1.100000023841858, + minimum: -1.0, + maximum: 5.0, + absoluteSum: 19.0, + positionChecksum: 10.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, -1.0, -1.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/index") + func test_uniqued_returning_index() throws { + try withIntegrationState(seed: 59710) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .index).index! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .uint32, + mean: 3.9000000953674316, + minimum: 0.0, + maximum: 10.0, + absoluteSum: 39.0, + positionChecksum: 20.2, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 7.0, 0.0, 10.0, 3.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/index/truncated") + func test_uniqued_returning_index_truncated() throws { + try withIntegrationState(seed: 97801) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .index).index! + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .uint32, + mean: 2.0, + minimum: 1.0, + maximum: 3.0, + absoluteSum: 6.0, + positionChecksum: 4.333333333333333, + sampleIndices: [0, 1, 2], + samples: [2.0, 1.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverse") + func test_uniqued_returning_inverse() throws { + try withIntegrationState(seed: 34157) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .inverse).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.899999976158142, + minimum: 0.0, + maximum: 4.0, + absoluteSum: 38.0, + positionChecksum: 22.35, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [2.0, 1.0, 0.0, 4.0, 2.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverse/truncated") + func test_uniqued_returning_inverse_truncated() throws { + try withIntegrationState(seed: 75343) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .inverse).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.5, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 30.0, + positionChecksum: 15.55, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [0.0, 1.0, 2.0, 2.0, 2.0, 2.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/counts") + func test_uniqued_returning_counts() throws { + try withIntegrationState(seed: 13539) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: .counts).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 6.0, + absoluteSum: 20.0, + positionChecksum: 6.9, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 5.0, 5.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/counts/truncated") + func test_uniqued_returning_counts_truncated() throws { + try withIntegrationState(seed: 70850) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 3, returning: .counts).counts! + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .int32, + mean: 4.0, + minimum: 2.0, + maximum: 6.0, + absoluteSum: 12.0, + positionChecksum: 8.666666666666666, + sampleIndices: [0, 1, 2], + samples: [4.0, 2.0, 6.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/index") + func test_uniqued_returning_all_index() throws { + try withIntegrationState(seed: 80422) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).index! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .uint32, + mean: 6.700000286102295, + minimum: 0.0, + maximum: 12.0, + absoluteSum: 67.0, + positionChecksum: 37.1, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [8.0, 0.0, 3.0, 1.0, 8.0, 8.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/inverse") + func test_uniqued_returning_all_inverse() throws { + try withIntegrationState(seed: 44680) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 1.600000023841858, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 32.0, + positionChecksum: 14.9, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [1.0, 2.0, 2.0, 1.0, 1.0, 3.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/all/counts") + func test_uniqued_returning_all_counts() throws { + try withIntegrationState(seed: 302) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 9.0, + absoluteSum: 20.0, + positionChecksum: 6.2, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [2.0, 9.0, 3.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/indexCounts") + func test_uniqued_returning_indexCounts() throws { + // skips .inverse: counts must be read from slot 2, not 3 + try withIntegrationState(seed: 95441) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.index, .counts]).counts! + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 2.0, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 20.0, + positionChecksum: 7.4, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [3.0, 5.0, 2.0, 4.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("uniqued/returning/inverseCounts") + func test_uniqued_returning_inverseCounts() throws { + try withIntegrationState(seed: 29675) { + let a = MLXRandom.randInt(low: 0, high: 6, [4, 5], type: Int32.self) + let result = MLX.uniqued(a, size: 10, returning: [.inverse, .counts]).inverse! + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .uint32, + mean: 2.6500000953674316, + minimum: 0.0, + maximum: 5.0, + absoluteSum: 53.0, + positionChecksum: 28.6, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [5.0, 0.0, 5.0, 2.0, 5.0, 4.0]), + tolerance: .exact) + } + } } diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift index 4fd0b1f9..9df64346 100644 --- a/Tests/MLXTests/FP8Tests.swift +++ b/Tests/MLXTests/FP8Tests.swift @@ -9,13 +9,26 @@ class FP8Tests: XCTestCase { override class func setUp() { } - func testFromFP8DecodesE4M3Bytes() { - let allBytes = (0 ... 255).map { UInt8($0) } + func testFromFP8DecodesE4M3BytesPositive() { + // 127 and 255 are NaN and don't compare equal, so skip them + let allBytes = (0 ... 126).map { UInt8($0) } let bytes = MLXArray(allBytes) let decoded = fromFP8(bytes, dtype: .float32) XCTAssertEqual(decoded.dtype, .float32) - XCTAssertEqual(decoded.shape, [256]) + XCTAssertEqual(decoded.shape, [127]) + + let expected = MLXArray(allBytes.map { e4m3Value($0) }) + assertEqual(decoded, expected) + } + + func testFromFP8DecodesE4M3BytesNegative() { + let allBytes = (128 ... 254).map { UInt8($0) } + let bytes = MLXArray(allBytes) + let decoded = fromFP8(bytes, dtype: .float32) + + XCTAssertEqual(decoded.dtype, .float32) + XCTAssertEqual(decoded.shape, [127]) let expected = MLXArray(allBytes.map { e4m3Value($0) }) assertEqual(decoded, expected) @@ -38,6 +51,7 @@ class FP8Tests: XCTestCase { } private func e4m3Value(_ value: UInt8) -> Float { + if (value & 0x7F) == 0x7F { return .nan } let sign: Float = (value & 0x80) == 0 ? 1 : -1 let exponent = Int((value >> 3) & 0x0F) let mantissa = Int(value & 0x07) diff --git a/Tests/MLXTests/MemoryTests.swift b/Tests/MLXTests/MemoryTests.swift index 42485e76..dc16ca9b 100644 --- a/Tests/MLXTests/MemoryTests.swift +++ b/Tests/MLXTests/MemoryTests.swift @@ -27,4 +27,54 @@ class MemoryTests: XCTestCase { XCTAssertEqual(again.maxBufferSize, info.maxBufferSize) XCTAssertEqual(again.maxRecommendedWorkingSetSize, info.maxRecommendedWorkingSetSize) } + + // `Memory.bufferSize` reports the allocator length of each buffer, and the + // allocator may hand out a cached buffer up to ~2x the requested size, so the + // exact value depends on the allocation history of the process (and other + // tests running in parallel). These tests only check the properties that hold + // regardless of the cache. + + func testBufferSizeAtLeastLogicalSize() { + let a = MLXRandom.normal([64, 32]) + let b = MLXArray.zeros([10, 10], type: UInt8.self) + eval(a, b) + + XCTAssertGreaterThanOrEqual(Memory.bufferSize([a]), a.nbytes) + XCTAssertGreaterThanOrEqual(Memory.bufferSize([b]), b.nbytes) + } + + func testBufferSizeEmpty() { + XCTAssertEqual(Memory.bufferSize([MLXArray]()), 0) + } + + func testBufferSizeCountsDuplicatesOnce() { + let a = MLXRandom.normal([64, 32]) + eval(a) + + XCTAssertEqual(Memory.bufferSize([a, a, a]), Memory.bufferSize([a])) + } + + func testBufferSizeCountsViewsOnce() { + let a = MLXRandom.normal([64, 32]) + let slice = a[0 ..< 8] + let transposed = a.T + eval(a, slice, transposed) + + // views share the buffer of `a` + let size = Memory.bufferSize([a]) + XCTAssertEqual(Memory.bufferSize([a, slice]), size) + XCTAssertEqual(Memory.bufferSize([a, transposed]), size) + XCTAssertEqual(Memory.bufferSize([slice, transposed]), size) + } + + func testBufferSizeSumsDistinctBuffers() { + let a = MLXRandom.normal([64, 32]) + let b = MLXRandom.normal([5]) + let c = MLXRandom.randInt(low: 0, high: 4, [100], type: UInt32.self) + eval(a, b, c) + + XCTAssertEqual( + Memory.bufferSize([a, b, c]), + Memory.bufferSize([a]) + Memory.bufferSize([b]) + Memory.bufferSize([c])) + } } diff --git a/cmake/mlx.patch b/cmake/mlx.patch deleted file mode 100644 index de99b7eb..00000000 --- a/cmake/mlx.patch +++ /dev/null @@ -1,34 +0,0 @@ -diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h -index f8a234ee6..3b7a79f27 100644 ---- a/mlx/backend/cuda/cuda_utils.h -+++ b/mlx/backend/cuda/cuda_utils.h -@@ -30,7 +30,11 @@ class CudaHandle { - if (cudaPeekAtLastError() != cudaSuccess) { - return; - } -- reset(); -+ // Inline reset() without throwing -- handle shutdown cases. -+ if (handle_ != nullptr) { -+ Destroy(handle_); -+ handle_ = nullptr; -+ } - } - - CudaHandle(const CudaHandle&) = delete; -diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp -index 472b3d99f..e7d8f2620 100644 ---- a/mlx/backend/cuda/device.cpp -+++ b/mlx/backend/cuda/device.cpp -@@ -617,8 +617,10 @@ std::unordered_map& get_command_encoders() { - } - - std::unordered_map& get_global_command_encoders() { -- static std::unordered_map encoders; -- return encoders; -+ // encoders are leaked intentionally as they would synchronize on process -+ // shutdown -+ static auto* encoders = new std::unordered_map(); -+ return *encoders; - } - - } // namespace mlx::core::cu diff --git a/tools/fix-metal-includes.sh b/tools/fix-metal-includes.sh index 6606683c..c32b8201 100755 --- a/tools/fix-metal-includes.sh +++ b/tools/fix-metal-includes.sh @@ -23,12 +23,16 @@ KERNEL_LIST=" \ arg_reduce.metal \ conv.metal \ dot.metal \ +gather_mm_offsets.metal \ layer_norm.metal \ random.metal \ rms_norm.metal \ rope.metal \ searchsorted.metal \ scaled_dot_product_attention.metal \ +scaled_dot_product_attention_vjp.metal \ +gated_delta_update.metal \ +gated_delta_update_vjp.metal \ steel/attn/kernels/steel_attention.metal" # We fixup all the header files AND the listed kernel files diff --git a/tools/integration_tests/README.md b/tools/integration_tests/README.md index 77eec885..567fc2f2 100644 --- a/tools/integration_tests/README.md +++ b/tools/integration_tests/README.md @@ -35,14 +35,14 @@ pair (`MAINTENANCE.md` step 10). | `Elementwise` | unary math, predicates, dtype conversions, complex parts | | `Binary` | binary ops, operators, comparisons, logical/bitwise, matmul family | | `Reduction` | reductions, cumulative ops, argMin/argMax, softmax | -| `Shape` | shape/structure ops, sort/partition, take/put, trace, hadamard | +| `Shape` | shape/structure ops, sort/partition, take/put, trace, hadamard, `matrixTransposed`, `uniqued` | | `Indexing` | subscript forms (ranges, `.ellipsis`, `.newAxis`, strides, arrays) | | `Factory` | `arange`, `linspace`, `zeros`/`ones`/`full`, `eye`, windows | | `Convolution` | `convolve`, `conv1d/2d/3d`, `convTransposed*`, `convGeneral` | -| `Quantization` | `quantized`/`dequantized`/`quantizedMM` per bits, gather/masked MM, fp8 | +| `Quantization` | `quantized`/`dequantized`/`quantizedMM` per bits, gather/masked MM, fp8, nvfp4 `globalScale` | | `FFT` | the whole `mx.fft` surface, including complex results | | `Linalg` | single-value linalg (`inv`, `cholesky`, `solve`, `det`, ...) | -| `Fast` | `rmsNorm`, `layerNorm`, `RoPE`, `scaledDotProductAttention` | +| `Fast` | `rmsNorm`, `layerNorm`, `RoPE`, `scaledDotProductAttention`, `gatedDeltaUpdate` | | `Random` | `gumbel`, `laplace`, `truncatedNormal`, `categorical`, `permutation`, ... | | `Activations` | the `MLXNN` activation free functions | | `Losses` | the `MLXNN` loss functions, per `reduction` | @@ -381,6 +381,11 @@ and exits before generating. bugs -- `MLX.nanToNum` defaulted `posInf`/`negInf` to `0` where python uses the dtype max/min (contradicting its own doc comment), and `MLX.tensordot` defaulted `axes` to `1` where python/numpy use `2`. Both are fixed. +- `Memory.bufferSize` / `mx.get_array_buffer_size` cannot be compared against + python: it reports the allocator length of each buffer, and the buffer cache + reuses any cached buffer smaller than `min(2 * size, size + 2 * page)`, so the + value depends on the allocation history of the process. It has property tests + in `Tests/MLXTests/MemoryTests.swift` instead. - `MLXRandom.laplace` requires `dtype:` where python defaults it, so the case spells it out. - `linspace` used to take its dtype from the literal type (`Int` -> `int64`, diff --git a/tools/integration_tests/cases.py b/tools/integration_tests/cases.py index be7ad47d..b7beb6d0 100644 --- a/tools/integration_tests/cases.py +++ b/tools/integration_tests/cases.py @@ -977,6 +977,132 @@ def binary(name: str, py_name: str, swift_name: str, lhs=None, rhs=None) -> None "mx.hadamard_transform(a, scale=0.25)", "MLX.hadamardTransform(a, scale: 0.25)", ) +case( + "matrixTransposed", + "Shape", + {"a": Normal(S)}, + "mx.matrix_transpose(a)", + "MLX.matrixTransposed(a)", +) +case( + "matrixTransposed/4d", + "Shape", + {"a": Normal((2, 3, 4, 5))}, + "mx.matrix_transpose(a)", + "MLX.matrixTransposed(a)", + note="only the last two axes move, unlike transposed(a) which reverses all of them", +) + +# `uniqued` has a fixed output `size`; inputs draw from a small range so there +# are repeats (20 draws from 6 values). Multiple return values are compared +# one field at a time; the combined `returning:` cases check that the fields +# are read back in the right order. +UNIQUE_A = RandInt(0, 6, (4, 5)) +case( + "uniqued", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=6)", + "MLX.uniqued(a, size: 6)", +) +case( + "uniqued/padded", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10)", + "MLX.uniqued(a, size: 10)", + note="no fillValue: pins down the default padding", +) +case( + "uniqued/fillValue", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, fill_value=mx.array(-1))", + "MLX.uniqued(a, size: 10, fillValue: MLXArray(-1))", +) +case( + "uniqued/truncated", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=3)", + "MLX.uniqued(a, size: 3)", + note="keeps the smallest `size` unique elements", +) +case( + "uniqued/float", + "Shape", + {"a": Normal(S)}, + "mx.unique(a, size=12)", + "MLX.uniqued(a, size: 12)", +) +case( + "uniqued/bool", + "Shape", + {"a": Bernoulli(S)}, + "mx.unique(a, size=2, stream=mx.cpu)", + "MLX.uniqued(a, size: 2, stream: .cpu)", + note="bool needs the cpu stream (builds on sort)", +) +case( + "uniqued/returning/values", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, " + "return_counts=True, fill_value=mx.array(-1))[0]", + "MLX.uniqued(a, size: 10, fillValue: MLXArray(-1), " + "returning: [.index, .inverse, .counts]).values", +) +for _field, _flag in (("index", "return_index"), ("inverse", "return_inverse"), ("counts", "return_counts")): + case( + f"uniqued/returning/{_field}", + "Shape", + {"a": UNIQUE_A}, + f"mx.unique(a, size=10, {_flag}=True)[1]", + f"MLX.uniqued(a, size: 10, returning: .{_field}).{_field}!", + ) + case( + f"uniqued/returning/{_field}/truncated", + "Shape", + {"a": UNIQUE_A}, + f"mx.unique(a, size=3, {_flag}=True)[1]", + f"MLX.uniqued(a, size: 3, returning: .{_field}).{_field}!", + ) +case( + "uniqued/returning/all/index", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[1]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).index!", +) +case( + "uniqued/returning/all/inverse", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[2]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).inverse!", +) +case( + "uniqued/returning/all/counts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_inverse=True, return_counts=True)[3]", + "MLX.uniqued(a, size: 10, returning: [.index, .inverse, .counts]).counts!", +) +case( + "uniqued/returning/indexCounts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_index=True, return_counts=True)[2]", + "MLX.uniqued(a, size: 10, returning: [.index, .counts]).counts!", + note="skips .inverse: counts must be read from slot 2, not 3", +) +case( + "uniqued/returning/inverseCounts", + "Shape", + {"a": UNIQUE_A}, + "mx.unique(a, size=10, return_inverse=True, return_counts=True)[1]", + "MLX.uniqued(a, size: 10, returning: [.inverse, .counts]).inverse!", +) # ------------------------------------------------------------------- factory # @@ -1335,6 +1461,63 @@ def binary(name: str, py_name: str, swift_name: str, lhs=None, rhs=None) -> None "MLX.fromFP8(encoded, dtype: .float32)", ) +# nvfp4 with one global scale per expert: `quantize` takes a single scale, so +# each expert is quantized on its own and the results stacked. The experts get +# different magnitudes so a scale applied to the wrong expert is visible. +NVFP4_EXPERTS = { + "w": Expression( + py="mx.random.normal((4, 32, 64)) * mx.array([1.0, 2.0, 4.0, 8.0]).reshape(4, 1, 1)", + swift="MLXRandom.normal([4, 32, 64]) * MLXArray([1.0, 2.0, 4.0, 8.0] as [Float]).reshaped(4, 1, 1)", + ), + "globalScale": Expression( + py="mx.max(mx.abs(w), axis=(1, 2))", + swift="MLX.abs(w).max(axes: [1, 2])", + ), + "q": Expression( + py='[mx.quantize(w[e], mode="nvfp4", global_scale=globalScale[e]) for e in range(4)]', + swift="(0 ..< 4).map { MLX.quantized(w[$0], mode: .nvfp4, globalScale: globalScale[$0]) }", + array=False, + ), + "wq": Expression(py="mx.stack([p[0] for p in q])", swift="MLX.stacked(q.map { $0.wq })"), + "scales": Expression( + py="mx.stack([p[1] for p in q])", swift="MLX.stacked(q.map { $0.scales })" + ), + "x": Normal((3, 2, 64)), + "indices": Expression(py="mx.array([2, 0, 3], dtype=mx.uint32)", swift="MLXArray([2, 0, 3] as [UInt32])"), +} +case( + "gatherQuantizedMM/nvfp4/globalScale", + "Quantization", + NVFP4_EXPERTS, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, mode="nvfp4", global_scale=globalScale)', + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, " + "globalScale: globalScale)", + note="global scale is gpu only and needs indices", +) +case( + "gatherQuantizedMM/nvfp4/globalScale/transpose", + "Quantization", + {**NVFP4_EXPERTS, "x": Normal((3, 2, 32))}, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, transpose=False, mode="nvfp4", ' + "global_scale=globalScale)", + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, transpose: false, " + "mode: .nvfp4, globalScale: globalScale)", +) +case( + "gatherQuantizedMM/nvfp4/globalScale/sorted", + "Quantization", + { + **NVFP4_EXPERTS, + "indices": Expression( + py="mx.array([0, 2, 2], dtype=mx.uint32)", swift="MLXArray([0, 2, 2] as [UInt32])" + ), + }, + 'mx.gather_qmm(x, wq, scales, rhs_indices=indices, mode="nvfp4", global_scale=globalScale, ' + "sorted_indices=True)", + "MLX.gatherQuantizedMM(x, wq, scales: scales, rhsIndices: indices, mode: .nvfp4, " + "globalScale: globalScale, sortedIndices: true)", +) + # ----------------------------------------------------------------------- fft FFT_1D = Normal((100,)) @@ -1643,6 +1826,78 @@ def fft_case(name: str, py: str, swift: str, *, complex_input=True, shape=(100,) "scale: 0.35, mask: mask)", ) + +def gated_delta_inputs(B, T, Hk, Hv, Dk, Dv, *, initial_state=True, mask=None) -> dict: + """inputs for `gated_delta_update`; small `keys` keep the delta rule stable""" + inputs = { + "queries": Normal((B, T, Hk, Dk)), + "keys": Uniform((B, T, Hk, Dk), low=-0.1, high=0.1), + "values": Normal((B, T, Hv, Dv)), + "gamma": Uniform((B, T, Hv), low=0.5, high=1.0), + "beta": Uniform((B, T, Hv), low=0.1, high=0.9), + } + if initial_state: + inputs["initialState"] = Normal((B, Hv, Dv, Dk)) + if mask is not None: + inputs["mask"] = mask + return inputs + + +# (name, dims, extra inputs, python kwargs, swift arguments, note) +GATED_DELTA = [ + # the generic (ops) path: any shape the metal kernel does not support + ("fallback", (2, 5, 2, 4, 8, 8), {}, "", "", "Hv != Hk, shape not supported by the kernel"), + ( + "fallback/initialState", + (2, 5, 2, 4, 8, 8), + {"initial_state": True}, + ", initial_state=initialState", + ", initialState: initialState", + None, + ), + ( + "fallback/mask", + (2, 5, 2, 4, 8, 8), + { + "initial_state": True, + "mask": Expression( + py="mx.array([[True] * 5, [True, True, True, False, False]])", + swift="MLXArray([true, true, true, true, true, true, true, true, false, false], [2, 5])", + ), + }, + ", initial_state=initialState, mask=mask", + ", initialState: initialState, mask: mask", + "masked steps leave the state unchanged and output zeros", + ), + # Dk == Dv == 128 with these head counts runs the metal kernel + ( + "kernel", + (1, 3, 16, 32, 128, 128), + {"initial_state": True}, + ", initial_state=initialState", + ", initialState: initialState", + "a shape the metal kernel supports", + ), +] +for _name, _dims, _extra, _py_kwargs, _swift_args, _note in GATED_DELTA: + _extra = {"initial_state": False, **_extra} + for _index, _output in ((0, "output"), (1, "state")): # python index, Swift label + case( + f"gatedDeltaUpdate/{_name}/{_output}", + "Fast", + gated_delta_inputs(*_dims, **_extra), + f"mx.fast.gated_delta_update(queries, keys, values, gamma, beta{_py_kwargs})[{_index}]", + "MLXFast.gatedDeltaUpdate(queries: queries, keys: keys, values: values, " + f"gamma: gamma, beta: beta{_swift_args}).{_output}", + note=_note, + ) + +# `Memory.bufferSize` is deliberately not here: it reports the allocator length +# of each buffer, and the allocator hands out cached buffers up to ~2x the +# requested size, so the value depends on what the process freed earlier (the +# generator runs every case in one warm process). It is covered by property +# tests in Tests/MLXTests/MemoryTests.swift instead. + # -------------------------------------------------------------------- random case("gumbel", "Random", {}, "mx.random.gumbel([4, 3])", "MLXRandom.gumbel([4, 3])") diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 710a87ff..cf74e4e6 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -357,6 +357,7 @@ "examples/example-metal-kernel.c", "examples/example-ops.c", "examples/example-safe-tensors.c", + "examples/example-stream-tls.c", "examples/example-thread-stream.c", examples/example.c, LICENSE, @@ -480,13 +481,16 @@ benchmarks/python/distributed_bench.py, benchmarks/python/einsum_bench.py, benchmarks/python/fft_bench.py, + benchmarks/python/gated_delta_bench.py, benchmarks/python/gather_bench.py, benchmarks/python/gather_mm_bench.py, benchmarks/python/gather_qmm_bench.py, + benchmarks/python/gather_qqmm_bench.py, benchmarks/python/hadamard_bench.py, benchmarks/python/large_gemm_bench.py, benchmarks/python/layer_norm_bench.py, benchmarks/python/masked_scatter.py, + benchmarks/python/qqmm_bench.py, benchmarks/python/rms_norm_bench.py, benchmarks/python/rope_bench.py, benchmarks/python/scatter_bench.py, @@ -495,6 +499,8 @@ benchmarks/python/segmented_mm_bench.py, benchmarks/python/single_ops.py, benchmarks/python/slice_update_bench.py, + benchmarks/python/swiglu_bench.py, + benchmarks/python/swiglu_qmm_bench.py, benchmarks/python/synchronize_bench.py, benchmarks/python/time_utils.py, CITATION.cff, @@ -589,6 +595,16 @@ docs/src/usage/saving_and_loading.rst, docs/src/usage/unified_memory.rst, docs/src/usage/using_streams.rst, + examples/cmake_extension/axpby/axpby.cpp, + examples/cmake_extension/axpby/axpby.h, + examples/cmake_extension/axpby/axpby.metal, + examples/cmake_extension/bindings.cpp, + examples/cmake_extension/CMakeLists.txt, + examples/cmake_extension/pyproject.toml, + examples/cmake_extension/README.md, + examples/cmake_extension/requirements.txt, + examples/cmake_extension/setup.py, + examples/cmake_extension/test.py, examples/cmake_project/CMakeLists.txt, examples/cmake_project/example.cpp, examples/cmake_project/README.md, @@ -605,17 +621,13 @@ examples/export/README.md, examples/export/train_mlp.cpp, examples/export/train_mlp.py, - examples/extensions/axpby/axpby.cpp, - examples/extensions/axpby/axpby.h, - examples/extensions/axpby/axpby.metal, - examples/extensions/bindings.cpp, - examples/extensions/CMakeLists.txt, - examples/extensions/mlx_sample_extensions/__init__.py, - examples/extensions/pyproject.toml, - examples/extensions/README.md, - examples/extensions/requirements.txt, - examples/extensions/setup.py, - examples/extensions/test.py, + examples/metal_extension/activation.cpp, + examples/metal_extension/pyproject.toml, + examples/metal_extension/README.md, + examples/metal_extension/requirements.txt, + examples/metal_extension/setup.py, + examples/metal_extension/silu_and_mul.metal, + examples/metal_extension/test.py, examples/python/distributed_data_parallel.py, examples/python/distributed_tensor_parallel.py, examples/python/linear_regression.py, @@ -689,6 +701,7 @@ mlx/backend/cuda/binary/CMakeLists.txt, mlx/backend/cuda/binary/divide.cu, mlx/backend/cuda/binary/equal.cu, + mlx/backend/cuda/binary/floor_divide.cu, mlx/backend/cuda/binary/greater_equal.cu, mlx/backend/cuda/binary/greater.cu, mlx/backend/cuda/binary/less_equal.cu, @@ -703,6 +716,7 @@ mlx/backend/cuda/binary/power.cu, mlx/backend/cuda/binary/remainder.cu, mlx/backend/cuda/binary/subtract.cu, + mlx/backend/cuda/cholesky.cu, mlx/backend/cuda/CMakeLists.txt, mlx/backend/cuda/compiled.cpp, mlx/backend/cuda/conv.cpp, @@ -722,6 +736,8 @@ mlx/backend/cuda/cuda.h, mlx/backend/cuda/cudnn_utils.cpp, mlx/backend/cuda/cudnn_utils.h, + mlx/backend/cuda/cusolver_utils.cpp, + mlx/backend/cuda/cusolver_utils.h, mlx/backend/cuda/custom_kernel.cpp, mlx/backend/cuda/cutlass_utils.cuh, mlx/backend/cuda/delayload.cpp, @@ -867,6 +883,8 @@ mlx/backend/cuda/utils.cpp, mlx/backend/cuda/utils.h, mlx/backend/cuda/vector_types.cuh, + mlx/backend/cuda/wddm.cpp, + mlx/backend/cuda/wddm.h, mlx/backend/cuda/worker.cpp, mlx/backend/cuda/worker.h, mlx/backend/gpu/CMakeLists.txt, @@ -900,6 +918,8 @@ mlx/backend/metal/kernels/conv.metal, mlx/backend/metal/kernels/copy.h, mlx/backend/metal/kernels/copy.metal, + mlx/backend/metal/kernels/cross_entropy.h, + mlx/backend/metal/kernels/cross_entropy.metal, mlx/backend/metal/kernels/defines.h, mlx/backend/metal/kernels/dot.h, mlx/backend/metal/kernels/dot.metal, @@ -916,6 +936,16 @@ mlx/backend/metal/kernels/fp_quantized.metal, mlx/backend/metal/kernels/fp4.h, mlx/backend/metal/kernels/fp8.h, + mlx/backend/metal/kernels/gated_delta_nax_ops.h, + mlx/backend/metal/kernels/gated_delta_update_nax_vjp.h, + mlx/backend/metal/kernels/gated_delta_update_nax_vjp.metal, + mlx/backend/metal/kernels/gated_delta_update_nax.h, + mlx/backend/metal/kernels/gated_delta_update_nax.metal, + mlx/backend/metal/kernels/gated_delta_update_vjp.h, + mlx/backend/metal/kernels/gated_delta_update_vjp.metal, + mlx/backend/metal/kernels/gated_delta_update.h, + mlx/backend/metal/kernels/gated_delta_update.metal, + mlx/backend/metal/kernels/gather_mm_offsets.metal, mlx/backend/metal/kernels/gemv_masked.h, mlx/backend/metal/kernels/gemv_masked.metal, mlx/backend/metal/kernels/gemv.h, @@ -948,9 +978,11 @@ mlx/backend/metal/kernels/reduction/reduce_row.h, mlx/backend/metal/kernels/rms_norm.metal, mlx/backend/metal/kernels/rope.metal, + mlx/backend/metal/kernels/scaled_dot_product_attention_vjp.metal, mlx/backend/metal/kernels/scaled_dot_product_attention.metal, mlx/backend/metal/kernels/scan.h, mlx/backend/metal/kernels/scan.metal, + mlx/backend/metal/kernels/sdpa_blocked.h, mlx/backend/metal/kernels/sdpa_vector.h, mlx/backend/metal/kernels/searchsorted.h, mlx/backend/metal/kernels/searchsorted.metal, @@ -1001,6 +1033,8 @@ mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.h, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_thin_nax.metal, mlx/backend/metal/kernels/steel/gemm/loader.h, mlx/backend/metal/kernels/steel/gemm/mma.h, mlx/backend/metal/kernels/steel/gemm/nax.h, @@ -1200,6 +1234,7 @@ python/tests/mpi_test_distributed.py, python/tests/nccl_test_distributed.py, python/tests/ring_test_distributed.py, + python/tests/run.py, python/tests/test_array.py, python/tests/test_autograd.py, python/tests/test_bf16.py, @@ -1213,6 +1248,8 @@ python/tests/test_einsum.py, python/tests/test_eval.py, python/tests/test_export_import.py, + python/tests/test_extension.py, + python/tests/test_fast_gated_delta.py, python/tests/test_fast_sdpa.py, python/tests/test_fast.py, python/tests/test_fft.py, @@ -1232,6 +1269,9 @@ python/tests/test_tree.py, python/tests/test_upsample.py, python/tests/test_vmap.py, + python/tests/test_zero_copy.py, + "python/tests/xp-tests/skips.txt", + "python/tests/xp-tests/xfail.txt", README.md, setup.py, tests/allocator_tests.cpp, @@ -1242,6 +1282,7 @@ tests/CMakeLists.txt, tests/compile_tests.cpp, tests/creations_tests.cpp, + tests/cuda_tests.cpp, tests/custom_vjp_tests.cpp, tests/device_tests.cpp, tests/einsum_tests.cpp, @@ -1354,6 +1395,7 @@ metal/cexpf.h, metal/complex.h, metal/copy.h, + metal/cross_entropy.h, metal/defines.h, metal/dot.h, metal/erf.h, @@ -1365,6 +1407,11 @@ metal/fp_quantized.h, metal/fp4.h, metal/fp8.h, + metal/gated_delta_nax_ops.h, + metal/gated_delta_update_nax_vjp.h, + metal/gated_delta_update_nax.h, + metal/gated_delta_update_vjp.h, + metal/gated_delta_update.h, metal/gemv_masked.h, metal/gemv.h, metal/hadamard.h, @@ -1388,6 +1435,7 @@ metal/reduction/reduce_init.h, metal/reduction/reduce_row.h, metal/scan.h, + metal/sdpa_blocked.h, metal/sdpa_vector.h, metal/searchsorted.h, metal/softmax.h, @@ -1421,6 +1469,7 @@ metal/steel/gemm/kernels/steel_gemm_segmented.h, metal/steel/gemm/kernels/steel_gemm_splitk_nax.h, metal/steel/gemm/kernels/steel_gemm_splitk.h, + metal/steel/gemm/kernels/steel_gemm_thin_nax.h, metal/steel/gemm/loader.h, metal/steel/gemm/mma.h, metal/steel/gemm/nax.h, From 392aa34400708e2e73d9ba3e99a68055786ad9f3 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:13:39 -0700 Subject: [PATCH 09/11] fix compatibility issue in update(parameters:) (#510) * fix compatibility issue in update(parameters:) - with the change in #418 the implementation changed such that the recursion was on an internal method - this restores the previous functionality --- CMakeLists.txt | 2 +- Source/Cmlx/mlx-c | 2 +- Source/MLXNN/Module.swift | 41 ++++++++++++++++++++++---------- Tests/MLXTests/ModuleTests.swift | 34 ++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 14 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index cf07480a..3ddeeb7c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -57,7 +57,7 @@ set(MLX_C_BUILD_EXAMPLES OFF) FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" - GIT_TAG "19d3c9c180dddf51be441b55e7f454861899bc93") + GIT_TAG "v0.8.0") FetchContent_MakeAvailable(mlx-c) # swift-numerics diff --git a/Source/Cmlx/mlx-c b/Source/Cmlx/mlx-c index 19d3c9c1..a26f2d58 160000 --- a/Source/Cmlx/mlx-c +++ b/Source/Cmlx/mlx-c @@ -1 +1 @@ -Subproject commit 19d3c9c180dddf51be441b55e7f454861899bc93 +Subproject commit a26f2d584a14ec619ae8db9a09be03becd47598a diff --git a/Source/MLXNN/Module.swift b/Source/MLXNN/Module.swift index a2c0a030..f1f7ea0a 100644 --- a/Source/MLXNN/Module.swift +++ b/Source/MLXNN/Module.swift @@ -509,20 +509,37 @@ open class Module { modulePath: [String] = [] ) throws -> Self { _checkNotSealed() - try update(parameters: parameters, verify: verify, path: path, modulePath: modulePath) { - m, k, a, v in - a._updateInternal(v) - } + try update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath, + mutate: nil) return self } + /// Applies `parameters` to this module and recurses into its children. + /// + /// With a `nil` `mutate`, arrays update in place and children recurse through the + /// open `update(parameters:verify:path:modulePath:)` so subclass overrides run. + /// A custom `mutate` (see `materialize()`) recurses directly and skips overrides. func update( parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [], modulePath: [String] = [], - mutate: (Module, String, MLXArray, MLXArray) -> Void + mutate: ((Module, String, MLXArray, MLXArray) -> Void)? ) throws { let modulePath = modulePath + [describeType(self)] + func updateChild( + _ module: Module, _ parameters: ModuleParameters, path: [String] + ) throws { + if let mutate { + try module.update( + parameters: parameters, verify: verify, path: path, + modulePath: modulePath, mutate: mutate) + } else { + try module.update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath) + } + } + func apply( key: String, path: [String], _ item: ModuleItem, _ value: NestedItem ) throws { @@ -542,7 +559,11 @@ open class Module { path: path, modules: modulePath, expectedShape: p.shape, actualShape: newArray.shape) } - mutate(self, key, p, newArray) + if let mutate { + mutate(self, key, p, newArray) + } else { + p._updateInternal(newArray) + } case (.value(.parameters), .none): if Self.parameterIsValid(key) { @@ -586,14 +607,10 @@ open class Module { } case (.value(.module(let module)), .dictionary(let values)): - try module.update( - parameters: NestedDictionary(values: values), verify: verify, path: path, - modulePath: modulePath, mutate: mutate) + try updateChild(module, NestedDictionary(values: values), path: path) case (.value(.module(let module)), .none): - try module.update( - parameters: NestedDictionary(), verify: verify, path: path, - modulePath: modulePath, mutate: mutate) + try updateChild(module, NestedDictionary(), path: path) case (.none, .none), (.value(.none), .none), (.value(.other(_)), .none): break diff --git a/Tests/MLXTests/ModuleTests.swift b/Tests/MLXTests/ModuleTests.swift index 07a771b3..c314addf 100644 --- a/Tests/MLXTests/ModuleTests.swift +++ b/Tests/MLXTests/ModuleTests.swift @@ -444,6 +444,40 @@ class ModuleTests: XCTestCase { XCTAssertEqual(collect["interior2.children.1"], "Linear") } + func testUpdateParametersCallsChildOverride() throws { + // a child's update(parameters:verify:path:modulePath:) override runs + // when the update starts at an ancestor + class RecordingLinear: Linear { + var updateCount = 0 + + override func update( + parameters: ModuleParameters, verify: VerifyUpdate, path: [String] = [], + modulePath: [String] = [] + ) throws -> Self { + updateCount += 1 + return try super.update( + parameters: parameters, verify: verify, path: path, modulePath: modulePath) + } + } + + class Parent: Module { + let child = RecordingLinear(2, 2) + let list = [RecordingLinear(2, 2)] + } + + let m = Parent() + let weight = MLXArray.ones([2, 2]) + m.update( + parameters: .unflattened([ + "child.weight": weight, "list.0.weight": weight, + ])) + + XCTAssertEqual(m.child.updateCount, 1) + XCTAssertEqual(m.list[0].updateCount, 1) + assertEqual(m.child.weight, weight) + assertEqual(m.list[0].weight, weight) + } + func testUpdateModuleSameType() throws { // set a single module with same type let m = newStructureModule() From dcbbf1a74b6e2fd7059cae97a7ea81c5ad16f615 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Tue, 27 Oct 2026 20:35:45 -0700 Subject: [PATCH 10/11] fix(manifest): use unsafeFlags for shorten-64-to-32 under tools 5.12 .disableWarning requires PackageDescription 6.2, but this manifest declares swift-tools-version 5.12, so the package failed to parse and every CI job was cancelled before any step ran. Use .unsafeFlags(["-Wno-shorten-64-to-32"]) instead, which 5.12 supports. Co-Authored-By: Claude Haiku 5.5 --- Package.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Package.swift b/Package.swift index f8e7d0c8..80f9d282 100644 --- a/Package.swift +++ b/Package.swift @@ -247,7 +247,7 @@ let cmlx = Target.target( .headerSearchPath("json/single_include/nlohmann"), .headerSearchPath("fmt/include"), .define("MLX_VERSION", to: "\"0.32.0\""), - .disableWarning("shorten-64-to-32"), + .unsafeFlags(["-Wno-shorten-64-to-32"]), ], linkerSettings: linkerSettings, plugins: cudaBuildPlugins, From e1f3202ab46815b49e4eb35b60ab6ed7377e6fd3 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Tue, 27 Oct 2026 21:47:01 -0700 Subject: [PATCH 11/11] fix(mlx-c): bind allocator::Buffer by value in prefault Upstream mlx changed array::buffer() to return allocator::Buffer by value. The fork's prefault() took a non-const reference to that temporary via const_cast, which fails to compile ("const_cast from rvalue to reference type"). Copy the buffer into a local and read its raw pointer from that. This broke the MLX Tests (SwiftPM metallib), Build xcode/MLX.xcodeproj and Downstream Integration Test jobs on dcbbf1a. Co-Authored-By: Claude Haiku 5.5 --- Source/Cmlx/mlx-c/mlx/c/fast.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Source/Cmlx/mlx-c/mlx/c/fast.cpp b/Source/Cmlx/mlx-c/mlx/c/fast.cpp index e8e141aa..e89488d2 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fast.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/fast.cpp @@ -14,7 +14,8 @@ namespace mlx::core { void prefault(const array& a) { if (!a.data_shared_ptr()) return; - const uint8_t* ptr = static_cast(const_cast(a.buffer()).raw_ptr()); + allocator::Buffer buf = a.buffer(); + const uint8_t* ptr = static_cast(buf.raw_ptr()); if (!ptr) return; // volatile to prevent the compiler from optimizing the read away