diff --git a/inference_lib/setup.cfg b/inference_lib/setup.cfg index 318c8fde..f13b6519 100644 --- a/inference_lib/setup.cfg +++ b/inference_lib/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = aqlm -version = 1.1.1 +version = 1.1.2dev author = AQLM paper authors author_email = vahe527887@yandex.ru description = Efficiently run models quantized with AQLM diff --git a/inference_lib/src/aqlm/inference_kernels/kernel_selector.py b/inference_lib/src/aqlm/inference_kernels/kernel_selector.py index b2cdaa1e..1ce1aed9 100644 --- a/inference_lib/src/aqlm/inference_kernels/kernel_selector.py +++ b/inference_lib/src/aqlm/inference_kernels/kernel_selector.py @@ -35,6 +35,17 @@ def get_forward_pass_kernel( from .cuda_kernel import CUDA_FOLDER return torch.ops.aqlm.code1x16_matmat + elif (optimize_for_training, codebooks.device.type, num_codebooks, codebook_size, out_group_size, in_group_size) == ( + False, + "mps", + 1, + 65536, + 1, + 8, + ): + from .mps_kernel import MPS_FOLDER + + return torch.ops.aqlm.code1x16_matmat_mps elif ( optimize_for_training, codebooks.device.type, diff --git a/inference_lib/src/aqlm/inference_kernels/mps_kernel.h b/inference_lib/src/aqlm/inference_kernels/mps_kernel.h new file mode 100644 index 00000000..327bfcaf --- /dev/null +++ b/inference_lib/src/aqlm/inference_kernels/mps_kernel.h @@ -0,0 +1,61 @@ +/* +See the LICENSE.txt file for this sample’s licensing information. + +Abstract: +The shader code for the custom operation. +*/ + +#pragma once + +// Defines the Metal soft shrink custom kernel. +static char *CUSTOM_KERNEL = R"MPS_AQLM( +#include +using namespace metal; + +template +kernel void Code1x16MatVec( + constant uint16_t* A [[buffer(0)]], + constant T* B [[buffer(1)]], + device T* C [[buffer(2)]], + constant T* codebook [[buffer(3)]], + constant int& prob_m [[buffer(4)]], + constant int& prob_k [[buffer(5)]], + uint index [[thread_position_in_grid]] +) { + uint num_codes = prob_k / 8; + constant uint16_t* codes_row = A + index * num_codes; + + float res = 0; + for (uint i = 0; i < num_codes; ++i) { + constant T* encoded_vector = codebook + static_cast(codes_row[i]) * 8; + for (uint j = 0; j < 8; ++j) { + res += static_cast(encoded_vector[j] * B[i * 8 + j]); + } + } + C[index] = res; +} + +template +[[host_name("aqlm_gemv_1x16_kernel_half")]] +kernel void Code1x16MatVec( + constant uint16_t* A [[buffer(0)]], + constant half* B [[buffer(1)]], + device half* C [[buffer(2)]], + constant half* codebook [[buffer(3)]], + constant int& prob_m [[buffer(4)]], + constant int& prob_k [[buffer(5)]], + uint index [[thread_position_in_grid]] +); + +template +[[host_name("aqlm_gemv_1x16_kernel_float")]] +kernel void Code1x16MatVec( + constant uint16_t* A [[buffer(0)]], + constant float* B [[buffer(1)]], + device float* C [[buffer(2)]], + constant float* codebook [[buffer(3)]], + constant int& prob_m [[buffer(4)]], + constant int& prob_k [[buffer(5)]], + uint index [[thread_position_in_grid]] +); +)MPS_AQLM"; diff --git a/inference_lib/src/aqlm/inference_kernels/mps_kernel.mm b/inference_lib/src/aqlm/inference_kernels/mps_kernel.mm new file mode 100644 index 00000000..28768e51 --- /dev/null +++ b/inference_lib/src/aqlm/inference_kernels/mps_kernel.mm @@ -0,0 +1,141 @@ +/* +See the LICENSE.txt file for this sample’s licensing information. + +Abstract: +The code that registers a PyTorch custom operation. +*/ + + +#include +#include "mps_kernel.h" + +#import +#import + +// Helper function to retrieve the `MTLBuffer` from a `torch::Tensor`. +static inline id getMTLBufferStorage(const torch::Tensor& tensor) { + return __builtin_bit_cast(id, tensor.storage().data()); +} + +void dispatchCode1x16Matvec( + const torch::Tensor& A, + const torch::Tensor& B, + torch::Tensor& C, + const torch::Tensor& codebook +) { + @autoreleasepool { + id device = MTLCreateSystemDefaultDevice(); + NSError *error = nil; + + // Set the number of threads equal to the number of rows. + int numThreads = C.size(-1); + + int prob_m = C.size(0); + int prob_k = B.size(0); + + // Load the custom soft shrink shader. + id customKernelLibrary = [device newLibraryWithSource:[NSString stringWithUTF8String:CUSTOM_KERNEL] + options:nil + error:&error]; + TORCH_CHECK(customKernelLibrary, "Failed to to create custom kernel library, error: ", error.localizedDescription.UTF8String); + + std::string kernel_name = std::string("aqlm_gemv_1x16_kernel_") + (B.scalar_type() == torch::kFloat ? "float" : "half"); + id customSoftShrinkFunction = [customKernelLibrary newFunctionWithName:[NSString stringWithUTF8String:kernel_name.c_str()]]; + TORCH_CHECK(customSoftShrinkFunction, "Failed to create function state object for ", kernel_name.c_str()); + + // Create a compute pipeline state object for the soft shrink kernel. + id softShrinkPSO = [device newComputePipelineStateWithFunction:customSoftShrinkFunction error:&error]; + TORCH_CHECK(softShrinkPSO, error.localizedDescription.UTF8String); + + // Get a reference to the command buffer for the MPS stream. + id commandBuffer = torch::mps::get_command_buffer(); + TORCH_CHECK(commandBuffer, "Failed to retrieve command buffer reference"); + + // Get a reference to the dispatch queue for the MPS stream, which encodes the synchronization with the CPU. + dispatch_queue_t serialQueue = torch::mps::get_dispatch_queue(); + + dispatch_sync(serialQueue, ^(){ + // Start a compute pass. + id computeEncoder = [commandBuffer computeCommandEncoder]; + TORCH_CHECK(computeEncoder, "Failed to create compute command encoder"); + + // Encode the pipeline state object and its parameters. + [computeEncoder setComputePipelineState:softShrinkPSO]; + [computeEncoder setBuffer:getMTLBufferStorage(A) offset:A.storage_offset() * A.element_size() atIndex:0]; + [computeEncoder setBuffer:getMTLBufferStorage(B) offset:B.storage_offset() * B.element_size() atIndex:1]; + [computeEncoder setBuffer:getMTLBufferStorage(C) offset:C.storage_offset() * C.element_size() atIndex:2]; + [computeEncoder setBuffer:getMTLBufferStorage(codebook) offset:codebook.storage_offset() * codebook.element_size() atIndex:3]; + [computeEncoder setBytes:&prob_m length:sizeof(int) atIndex:4]; + [computeEncoder setBytes:&prob_k length:sizeof(int) atIndex:5]; + + MTLSize gridSize = MTLSizeMake(numThreads, 1, 1); + + // Calculate a thread group size. + NSUInteger threadGroupSize = softShrinkPSO.maxTotalThreadsPerThreadgroup; + if (threadGroupSize > numThreads) { + threadGroupSize = numThreads; + } + MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1); + + // Encode the compute command. + [computeEncoder dispatchThreads:gridSize + threadsPerThreadgroup:threadgroupSize]; + + [computeEncoder endEncoding]; + + // Commit the work. + torch::mps::commit(); + }); + } +} + +// C++ op dispatching the Metal soft shrink shader. +torch::Tensor code1x16_matmat( + const torch::Tensor& input, + const torch::Tensor& codes, + const torch::Tensor& codebooks, + const torch::Tensor& scales, + const std::optional& bias +) { + // Check whether the input tensor resides on the MPS device and whether it's contiguous. + TORCH_CHECK(input.device().is_mps(), "input must be a MPS tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + + // Check the supported data types for soft shrink. + TORCH_CHECK(input.scalar_type() == torch::kFloat || input.scalar_type() == torch::kHalf, "Unsupported data type: ", input.scalar_type()); + + auto input_sizes = input.sizes(); + auto out_features = codes.size(0) * codebooks.size(2); + auto flat_input = input.reshape({-1, input.size(-1)}); + auto flat_output = torch::empty({flat_input.size(0), out_features}, + torch::TensorOptions() + .dtype(input.dtype()) + .device(input.device()) + ); + + for (int i = 0; i < flat_input.size(0); ++i) { + auto input_vec = flat_input.index({i}); + auto output_vec = flat_output.index({i}); + dispatchCode1x16Matvec( + codes.squeeze(2), + input_vec, + output_vec, + codebooks + ); + } + flat_output *= scales.flatten().unsqueeze(0); + if (bias.has_value()) { + flat_output += bias->unsqueeze(0); + } + + auto output_sizes = input_sizes.vec(); + output_sizes.pop_back(); + output_sizes.push_back(-1); + auto output = flat_output.reshape(output_sizes).clone(); + return output; +} + +// Create Python bindings for the Objective-C++ code. +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("code1x16_matmat", &code1x16_matmat); +} diff --git a/inference_lib/src/aqlm/inference_kernels/mps_kernel.py b/inference_lib/src/aqlm/inference_kernels/mps_kernel.py new file mode 100644 index 00000000..51271129 --- /dev/null +++ b/inference_lib/src/aqlm/inference_kernels/mps_kernel.py @@ -0,0 +1,23 @@ +import os +from typing import Optional + +import torch +from torch.utils.cpp_extension import load + +MPS_FOLDER = os.path.dirname(os.path.abspath(__file__)) +MPS_KERNEL = load( + name="codebook_mps", + sources=[os.path.join(MPS_FOLDER, "mps_kernel.mm")], + extra_cflags=['-std=c++17'], +) + +torch.library.define( + "aqlm::code1x16_matmat_mps", "(Tensor input, Tensor codes, Tensor codebooks, Tensor scales, Tensor bias) -> Tensor" +) + +torch.library.impl("aqlm::code1x16_matmat_mps", "default", MPS_KERNEL.code1x16_matmat) + + +@torch.library.impl_abstract("aqlm::code1x16_matmat_mps") +def code1x16_matmat_meta(input, codes, codebooks, scales, bias): + return torch.empty(input.shape[:-1] + (codes.shape[0],), device=input.device, dtype=input.dtype)