diff --git a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.cc b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.cc new file mode 100644 index 000000000..22a9a5329 --- /dev/null +++ b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.cc @@ -0,0 +1,37 @@ +#include "linked/torch/nvidia/ops/awq_marlin_repack/vllm.h" + +#include +#include +#include + +#include +#include + +namespace infini::ops::linked::torch::nvidia { + +at::Tensor VllmAwqMarlinRepack::Call(at::Tensor b_q_weight, int64_t size_k, + int64_t size_n, int64_t num_bits, + bool is_a_8bit) { + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "_C::awq_marlin_repack", ""); + c10::Stack stack; + stack.emplace_back(std::move(b_q_weight)); + stack.emplace_back(c10::SymInt{size_k}); + stack.emplace_back(c10::SymInt{size_n}); + stack.emplace_back(num_bits); + stack.emplace_back(is_a_8bit); + op.callBoxed(&stack); + + assert(stack.size() == 1 && + "`awq_marlin_repack` returned an unexpected number of values"); + return std::move(stack.front()).toTensor(); +} + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +template class TorchAwqMarlinRepack< + ::infini::ops::linked::torch::nvidia::VllmAwqMarlinRepack>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h new file mode 100644 index 000000000..a350622dc --- /dev/null +++ b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h @@ -0,0 +1,39 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_ + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/awq_marlin_repack.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmAwqMarlinRepack : C10 { + static at::Tensor Call(at::Tensor b_q_weight, int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit); +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchAwqMarlinRepack< + ::infini::ops::linked::torch::nvidia::VllmAwqMarlinRepack>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchAwqMarlinRepack< + linked::torch::nvidia::VllmAwqMarlinRepack> { + public: + using linked::torch::TorchAwqMarlinRepack< + linked::torch::nvidia::VllmAwqMarlinRepack>::TorchAwqMarlinRepack; + + using linked::torch::TorchAwqMarlinRepack< + linked::torch::nvidia::VllmAwqMarlinRepack>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_ diff --git a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.yaml b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.yaml new file mode 100644 index 000000000..cdac1c767 --- /dev/null +++ b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.yaml @@ -0,0 +1,5 @@ +library: vllm +operator_schema: >- + _C::awq_marlin_repack(Tensor b_q_weight, SymInt size_k, SymInt size_n, + int num_bits, bool is_a_8bit) -> Tensor +dispatch_key: CUDA diff --git a/src/linked/torch/ops/awq_marlin_repack.h b/src/linked/torch/ops/awq_marlin_repack.h new file mode 100644 index 000000000..ba0e54958 --- /dev/null +++ b/src/linked/torch/ops/awq_marlin_repack.h @@ -0,0 +1,40 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_ + +#include + +#include "base/awq_marlin_repack.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchAwqMarlinRepack : public ::infini::ops::AwqMarlinRepack { + public: + using ::infini::ops::AwqMarlinRepack::AwqMarlinRepack; + using ::infini::ops::AwqMarlinRepack::operator(); + + void operator()(const Tensor b_q_weight, const int64_t size_k, + const int64_t size_n, const int64_t num_bits, + const bool is_a_8bit, Tensor out) const override { + ValidateCallMetadata(b_q_weight, size_k, size_n, num_bits, is_a_8bit, out); + + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + auto at_b_q_weight = ToAtenTensor( + const_cast(b_q_weight.data()), b_q_weight_metadata_.shape(), + b_q_weight_metadata_.strides(), b_q_weight_metadata_.dtype(), + device_index_); + auto at_out = ToAtenTensor( + out.data(), out_metadata_.shape(), out_metadata_.strides(), + out_metadata_.dtype(), device_index_); + + auto result = Backend::Call(std::move(at_b_q_weight), size_k, size_n, + num_bits, is_a_8bit); + at_out.copy_(result); + } +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_ diff --git a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cu b/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cu deleted file mode 100644 index c4bb99211..000000000 --- a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cu +++ /dev/null @@ -1,109 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from vLLM at commit 25ace8fe5df07fc13f4aef5a89db391f326e60ee: -// csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu - -#include - -#include -#include - -#include "native/cuda/nvidia/ops/awq_marlin_repack/kernel.cuh" -#include "native/cuda/nvidia/ops/awq_marlin_repack/kernel.h" - -namespace infini::ops { -namespace { - -class DeviceGuard { - public: - explicit DeviceGuard(int device_index) { - auto status = cudaGetDevice(&previous_device_); - assert(status == cudaSuccess && - "`AwqMarlinRepack` failed to query the current CUDA device"); - - if (previous_device_ != device_index) { - status = cudaSetDevice(device_index); - assert(status == cudaSuccess && - "`AwqMarlinRepack` failed to select the input CUDA device"); - restore_ = true; - } - } - - ~DeviceGuard() { - if (restore_) { - const auto status = cudaSetDevice(previous_device_); - assert(status == cudaSuccess && - "`AwqMarlinRepack` failed to restore the CUDA device"); - } - } - - private: - int previous_device_{0}; - - bool restore_{false}; -}; - -template -void Launch(const uint32_t* b_q_weight, uint32_t* out, int size_k, int size_n, - int blocks, int shared_memory_bytes, cudaStream_t stream) { - const auto attribute_status = cudaFuncSetAttribute( - awq_marlin_repack_detail::AwqMarlinRepackKernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_bytes); - assert(attribute_status == cudaSuccess && - "`AwqMarlinRepack` failed to configure dynamic shared memory"); - - awq_marlin_repack_detail::AwqMarlinRepackKernel - <<>>(b_q_weight, out, size_k, size_n); -} - -} // namespace - -void Operator::operator()( - const Tensor b_q_weight, const int64_t size_k, const int64_t size_n, - const int64_t num_bits, const bool is_a_8bit, Tensor out) const { - ValidateCallMetadata(b_q_weight, size_k, size_n, num_bits, is_a_8bit, out); - - DeviceGuard device_guard{device_index_}; - int blocks = 0; - auto status = cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, - device_index_); - assert(status == cudaSuccess && blocks > 0 && - "`AwqMarlinRepack` failed to query CUDA multiprocessor count"); - - int shared_memory_bytes = 0; - status = cudaDeviceGetAttribute(&shared_memory_bytes, - cudaDevAttrMaxSharedMemoryPerBlockOptin, - device_index_); - assert(status == cudaSuccess && shared_memory_bytes > 0 && - "`AwqMarlinRepack` failed to query CUDA shared memory capacity"); - - const auto* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data()); - auto* out_ptr = reinterpret_cast(out.data()); - const auto stream = static_cast(stream_ ? stream_ : 0); - const auto kernel_size_k = static_cast(size_k_); - const auto kernel_size_n = static_cast(size_n_); - - if (is_a_8bit_) { - if (num_bits_ == 4) { - Launch<4, true>(b_q_weight_ptr, out_ptr, kernel_size_k, kernel_size_n, - blocks, shared_memory_bytes, stream); - } else { - Launch<8, true>(b_q_weight_ptr, out_ptr, kernel_size_k, kernel_size_n, - blocks, shared_memory_bytes, stream); - } - } else if (num_bits_ == 4) { - Launch<4, false>(b_q_weight_ptr, out_ptr, kernel_size_k, kernel_size_n, - blocks, shared_memory_bytes, stream); - } else { - Launch<8, false>(b_q_weight_ptr, out_ptr, kernel_size_k, kernel_size_n, - blocks, shared_memory_bytes, stream); - } - - status = cudaGetLastError(); - assert(status == cudaSuccess && - "`AwqMarlinRepack` CUDA kernel launch failed"); -} - -} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cuh b/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cuh deleted file mode 100644 index 34a159101..000000000 --- a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.cuh +++ /dev/null @@ -1,240 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from vLLM at commit 25ace8fe5df07fc13f4aef5a89db391f326e60ee: -// csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu -// csrc/libtorch_stable/quantization/marlin/marlin.cuh - -#ifndef INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_CUH_ -#define INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_CUH_ - -#include -#include - -#include - -namespace infini::ops::awq_marlin_repack_detail { - -constexpr int kRepackStages = 8; -constexpr int kRepackThreads = 256; -constexpr int kTileSize = 16; -constexpr int kTileKSize = kTileSize; -constexpr int kTileNSize = kTileKSize * 4; - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void CpAsync4(void* shared_ptr, const void* global_ptr) { - reinterpret_cast(shared_ptr)[0] = - reinterpret_cast(global_ptr)[0]; -} - -__device__ inline void CpAsyncFence() {} - -template -__device__ inline void CpAsyncWait() {} - -#else - -__device__ inline void CpAsync4(void* shared_ptr, const void* global_ptr) { - constexpr int kBytes = 16; - const auto shared_address = - static_cast(__cvta_generic_to_shared(shared_ptr)); - asm volatile("cp.async.cg.shared.global [%0], [%1], %2;\n" - : - : "r"(shared_address), "l"(global_ptr), "n"(kBytes)); -} - -__device__ inline void CpAsyncFence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void CpAsyncWait() { - asm volatile("cp.async.wait_group %0;\n" : : "n"(kCount)); -} - -#endif - -template -__global__ void AwqMarlinRepackKernel( - const uint32_t* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, - int size_k, int size_n) { - constexpr int kPackFactor = 32 / kNumBits; - constexpr int kTargetTileNSize = kTileNSize / (kIsA8Bit ? 2 : 1); - constexpr int kTargetTileKSize = kTileKSize * (kIsA8Bit ? 2 : 1); - const int k_tiles = size_k / kTargetTileKSize; - const int n_tiles = size_n / kTargetTileNSize; - const int block_k_tiles = (k_tiles + gridDim.x - 1) / gridDim.x; - - const auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - const int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - const auto wait_for_stage = [&]() { - CpAsyncWait(); - __syncthreads(); - }; - - extern __shared__ int4 shared[]; - constexpr int kTileNInts = kTargetTileNSize / kPackFactor; - constexpr int kStageNThreads = kTileNInts / 4; - constexpr int kStageKThreads = kTargetTileKSize; - constexpr int kStageSize = kStageKThreads * kStageNThreads; - - const auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - CpAsyncFence(); - return; - } - - const int first_n = n_tile_id * kTargetTileNSize; - const int first_n_packed = first_n / kPackFactor; - int4* shared_ptr = shared + kStageSize * pipe; - - if (threadIdx.x < kStageSize) { - const auto k_id = threadIdx.x / kStageNThreads; - const auto n_id = threadIdx.x % kStageNThreads; - const int first_k = k_tile_id * kTargetTileKSize; - - CpAsync4(&shared_ptr[k_id * kStageNThreads + n_id], - reinterpret_cast( - &b_q_weight_ptr[(first_k + k_id) * (size_n / kPackFactor) + - first_n_packed + n_id * 4])); - } - - CpAsyncFence(); - }; - - const auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - const auto warp_id = threadIdx.x / 32; - const auto thread_id = threadIdx.x % 32; - if (warp_id >= 4) { - return; - } - - const int tensor_core_column = thread_id / 4; - const int tensor_core_row = (thread_id % 4) * (kIsA8Bit ? 4 : 2); - constexpr int kTensorCoreOffsets[4] = {0, 1, 8, 9}; - const int current_n = - (warp_id / (kIsA8Bit ? 2 : 1)) * 16 + tensor_core_column; - const int current_n_packed = current_n / kPackFactor; - const int current_n_position = current_n % kPackFactor; - - constexpr int kSharedStride = kTileNInts; - constexpr uint32_t kMask = (1U << kNumBits) - 1; - int4* shared_stage_ptr = shared + kStageSize * pipe; - auto* shared_stage_int_ptr = reinterpret_cast(shared_stage_ptr); - - int unpacked_n_position = 0; - if constexpr (kNumBits == 4) { - constexpr int kUndoPack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - unpacked_n_position = kUndoPack[current_n_position]; - } else { - constexpr int kUndoPack[4] = {0, 2, 1, 3}; - unpacked_n_position = kUndoPack[current_n_position]; - } - - uint32_t values[8]; -#pragma unroll - for (int i = 0; i < 4; ++i) { - if constexpr (kIsA8Bit) { - const int current_element = tensor_core_row + i; - const uint32_t first_value = - shared_stage_int_ptr[current_n_packed + - (8 / kPackFactor) * (warp_id % 2) + - kSharedStride * current_element]; - const uint32_t second_value = - shared_stage_int_ptr[current_n_packed + - (8 / kPackFactor) * (warp_id % 2) + - kSharedStride * (current_element + 16)]; - - values[i] = (first_value >> (unpacked_n_position * kNumBits)) & kMask; - values[4 + i] = - (second_value >> (unpacked_n_position * kNumBits)) & kMask; - } else { - const int current_element = tensor_core_row + kTensorCoreOffsets[i]; - const uint32_t first_value = - shared_stage_int_ptr[current_n_packed + - kSharedStride * current_element]; - const uint32_t second_value = - shared_stage_int_ptr[current_n_packed + (8 / kPackFactor) + - kSharedStride * current_element]; - - values[i] = (first_value >> (unpacked_n_position * kNumBits)) & kMask; - values[4 + i] = - (second_value >> (unpacked_n_position * kNumBits)) & kMask; - } - } - - constexpr int kTileElements = - kTargetTileKSize * kTargetTileNSize / kPackFactor; - const int out_offset = (k_tile_id * n_tiles + n_tile_id) * kTileElements; - - // Matches FasterTransformer's interleaved numeric conversion layout: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!kIsA8Bit && kNumBits == 4) { - constexpr int kPackIndices[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - uint32_t result = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) { - result |= values[kPackIndices[i]] << (i * 4); - } - out_ptr[out_offset + thread_id * 4 + warp_id] = result; - } else if constexpr (kIsA8Bit && kNumBits == 4) { - constexpr int kPackIndices[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - uint32_t result = 0; -#pragma unroll - for (int i = 0; i < 8; ++i) { - result |= values[kPackIndices[i]] << (i * 4); - } - out_ptr[out_offset + thread_id * 4 + warp_id] = result; - } else { - constexpr int kPackIndices[4] = {0, 2, 1, 3}; - uint32_t first_result = 0; - uint32_t second_result = 0; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int index = kIsA8Bit ? i : kPackIndices[i]; - first_result |= values[index] << (i * 8); - second_result |= values[4 + index] << (i * 8); - } - out_ptr[out_offset + thread_id * 8 + warp_id * 2] = first_result; - out_ptr[out_offset + thread_id * 8 + warp_id * 2 + 1] = second_result; - } - }; - - const auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < kRepackStages - 1; ++pipe) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - wait_for_stage(); - }; - -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; ++k_tile_id) { - int n_tile_id = 0; - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < kRepackStages; ++pipe) { - fetch_to_shared((pipe + kRepackStages - 1) % kRepackStages, k_tile_id, - n_tile_id + pipe + kRepackStages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += kRepackStages; - } - } -} - -} // namespace infini::ops::awq_marlin_repack_detail - -#endif // INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_CUH_ diff --git a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.h b/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.h deleted file mode 100644 index faeb6984f..000000000 --- a/src/native/cuda/nvidia/ops/awq_marlin_repack/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_H_ -#define INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_H_ - -#include "base/awq_marlin_repack.h" - -namespace infini::ops { - -template <> -class Operator - : public AwqMarlinRepack { - public: - using AwqMarlinRepack::AwqMarlinRepack; - - void operator()(const Tensor b_q_weight, const int64_t size_k, - const int64_t size_n, const int64_t num_bits, - const bool is_a_8bit, Tensor out) const override; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_NVIDIA_AWQ_MARLIN_REPACK_KERNEL_H_