Skip to content

feat(linear): expose stream-aware W4A16 API #17

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
qescccczmr wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat(linear): expose stream-aware W4A16 API #17

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions src/turbomind/api/python/bind.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,15 +201,20 @@ std::shared_ptr<turbomind::Tensor> DLManagedTensorToTurbomindTensor(DLManagedTen
return std::make_shared<turbomind::Tensor>(where, dtype, shape, data);
}

std::shared_ptr<turbomind::Tensor> TorchTensorToTurbomindTensor(py::object obj)
std::shared_ptr<turbomind::Tensor> TorchTensorToTurbomindTensor(py::object obj, int64_t stream_id)
{
py::capsule cap = obj.attr("__dlpack__")();
// DLPack represents CUDA's legacy default stream as 1, while the native
// CUDA stream handle is nullptr (0).
const auto dlpack_stream = stream_id == 0 ? 1 : stream_id;
py::capsule cap = obj.attr("__dlpack__")("stream"_a = dlpack_stream);
DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(PyCapsule_GetPointer(cap.ptr(), kDlTensorCapsuleName));
return DLManagedTensorToTurbomindTensor(dlmt);
}

PYBIND11_MODULE(_turbomind_ext, m)
{
m.attr("W4A16_LINEAR_API_VERSION") = 2;

py::enum_<turbomind::WeightType>(m, "WeightType")
.value("kFP32", turbomind::WeightType::kFP32)
.value("kFP16", turbomind::WeightType::kFP16)
Expand Down Expand Up @@ -331,16 +336,22 @@ PYBIND11_MODULE(_turbomind_ext, m)
return new turbomind::Linear(in_features, out_features, w_bit, group_size);
}))
.def("post_init",
[](turbomind::Linear* self, py::object qweight, py::object scales, py::object qzeros, bool simt) {
auto _qweight = TorchTensorToTurbomindTensor(qweight);
auto _scales = TorchTensorToTurbomindTensor(scales);
auto _qzeros = TorchTensorToTurbomindTensor(qzeros);
self->post_init(_qweight, *_scales, *_qzeros, simt);
[](turbomind::Linear* self,
py::object qweight,
py::object scales,
py::object qzeros,
bool simt,
int64_t stream_id = 0) {
auto _qweight = TorchTensorToTurbomindTensor(qweight, stream_id);
auto _scales = TorchTensorToTurbomindTensor(scales, stream_id);
auto _qzeros = TorchTensorToTurbomindTensor(qzeros, stream_id);
auto stream = reinterpret_cast<cudaStream_t>(stream_id);
self->post_init(_qweight, *_scales, *_qzeros, simt, stream);
})
.def("forward",
[](turbomind::Linear* self, py::object in, py::object out, int64_t stream_id = 0) {
auto _in = TorchTensorToTurbomindTensor(in);
auto _out = TorchTensorToTurbomindTensor(out);
auto _in = TorchTensorToTurbomindTensor(in, stream_id);
auto _out = TorchTensorToTurbomindTensor(out, stream_id);
auto stream = reinterpret_cast<cudaStream_t>(stream_id);
return self->forward(*_in, *_out, stream);
})
Expand Down
138 changes: 97 additions & 41 deletions src/turbomind/api/python/linear.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,27 +63,67 @@ struct Linear::Impl {
Impl(size_t input_dims, size_t output_dims, int w_bit, int group_size):
input_dims_(input_dims), output_dims_(output_dims), w_bit_(w_bit), group_size_(group_size)
{
TM_CHECK(w_bit_ == 4);
TM_CHECK(group_size_ == 128);
TM_CHECK(input_dims_ % group_size_ == 0);
TM_CHECK(output_dims_ % 32 == 0);
}

~Impl()
~Impl() noexcept
{
check_cuda_error(cudaFree(scales_zeros_));
if (scales_zeros_) {
int old_device{};
if (cudaGetDevice(&old_device) != cudaSuccess) {
return;
}
if (device_id_ >= 0 && old_device != device_id_ && cudaSetDevice(device_id_) != cudaSuccess) {
return;
}
// cudaFree does not synchronize memory allocated by
// cudaMallocAsync. Destruction is outside the inference path.
if (cudaDeviceSynchronize() == cudaSuccess) {
cudaFree(scales_zeros_);
}
if (device_id_ >= 0 && old_device != device_id_) {
cudaSetDevice(old_device);
}
}
}

void post_init(std::shared_ptr<Tensor> qweight, const Tensor& scales, const Tensor& qzeros, bool simt)
void post_init(
std::shared_ptr<Tensor> qweight, const Tensor& scales, const Tensor& qzeros, bool simt, cudaStream_t stream)
{
TM_CHECK(!weight_ && !scales_zeros_);
TM_CHECK(qweight != nullptr);
TM_CHECK(qweight->where == MEMORY_GPU && scales.where == MEMORY_GPU && qzeros.where == MEMORY_GPU);
TM_CHECK(qweight->type == TYPE_INT32 && scales.type == TYPE_FP16 && qzeros.type == TYPE_FP16);

const auto scale_count = input_dims_ / group_size_ * output_dims_;
TM_CHECK(qweight->sizeBytes() == input_dims_ * output_dims_ / 2);
TM_CHECK(scales.size() == scale_count && qzeros.size() == scale_count);
check_cuda_error(cudaGetDevice(&device_id_));

const auto workspace_size = input_dims_ * output_dims_ * sizeof(uint16_t);
void* workspace{};
check_cuda_error(cudaMalloc((void**)&workspace, workspace_size));
check_cuda_error(cudaMallocAsync(&workspace, workspace_size, stream));

convert_qweight(workspace, qweight, input_dims_, output_dims_, simt);
convert_scales_zeros(workspace, scales, qzeros, input_dims_, output_dims_, group_size_, simt);
try {
convert_qweight(workspace, qweight, input_dims_, output_dims_, simt, stream);
convert_scales_zeros(workspace, scales, qzeros, input_dims_, output_dims_, group_size_, simt, stream);
}
catch (...) {
cudaFreeAsync(workspace, stream);
throw;
}

check_cuda_error(cudaFree(workspace));
check_cuda_error(cudaFreeAsync(workspace, stream));
weight_ = qweight;
}

void forward(const Tensor& in, Tensor& out, cudaStream_t stream)
{
TM_CHECK(weight_ && scales_zeros_);
TM_CHECK(in.where == MEMORY_GPU && out.where == MEMORY_GPU);
TM_CHECK(in.type == TYPE_FP16 && out.type == TYPE_FP16);
TM_CHECK(in.shape.size() == 2 && in.shape[1] == input_dims_);
TM_CHECK(out.shape.size() == 2 && out.shape[0] == in.shape[0] && out.shape[1] == output_dims_);
Expand All @@ -108,6 +148,7 @@ struct Linear::Impl {
(int)output_dims_};
int device_id;
check_cuda_error(cudaGetDevice(&device_id));
TM_CHECK(device_id == device_id_);
auto gemm = GemmPool::getInstance().get(device_id);
auto ec = gemm->Run(operation,
1.f,
Expand All @@ -128,23 +169,27 @@ struct Linear::Impl {
stream);

if (ec) {
printf("%s: %d", __PRETTY_FUNCTION__, ec);
std::abort();
TM_THROW("W4A16 GEMM failed with error code " + std::to_string(ec) + ".");
}
}
void convert_qweight(
void* workspace, std::shared_ptr<Tensor> weight, size_t input_dims, size_t output_dims, bool use_simt)
void convert_qweight(void* workspace,
std::shared_ptr<Tensor> weight,
size_t input_dims,
size_t output_dims,
bool use_simt,
cudaStream_t stream)
{
using namespace gemm;
auto [order_b, pack_b, order_v, pack_v] = get_weight_and_scales_layout(getSMVersion(), use_simt);

if (order_b == kColMajor) {
transpose_u4((uint4_t*)workspace, (const uint4_t*)weight->data, input_dims, output_dims);
cudaMemcpy(const_cast<void*>(weight->data), workspace, input_dims * output_dims / 2, cudaMemcpyDefault);
transpose_u4((uint4_t*)workspace, (const uint4_t*)weight->data, input_dims, output_dims, stream);
check_cuda_error(cudaMemcpyAsync(
const_cast<void*>(weight->data), workspace, input_dims * output_dims / 2, cudaMemcpyDefault, stream));
}

extend_to_u16((uint16_t*)workspace, (const uint4_t*)weight->data, input_dims * output_dims);
sync_check_cuda_error();
extend_to_u16((uint16_t*)workspace, (const uint4_t*)weight->data, input_dims * output_dims, stream);
check_cuda_error(cudaPeekAtLastError());

if constexpr (0) {
std::vector<uint16_t> tmp(input_dims * output_dims);
Expand Down Expand Up @@ -173,12 +218,10 @@ struct Linear::Impl {
k_desc_.type = gemm::DataType::U4;
k_desc_.pack = pack_b;

cudaMemset(const_cast<void*>(weight->data), 0, input_dims * output_dims / 2);
check_cuda_error(cudaMemsetAsync(const_cast<void*>(weight->data), 0, input_dims * output_dims / 2, stream));

TM_CHECK(Convert(workspace, w_desc, const_cast<void*>(weight->data), k_desc_, 0) == 0);
sync_check_cuda_error();

cudaDeviceSynchronize();
TM_CHECK(Convert(workspace, w_desc, const_cast<void*>(weight->data), k_desc_, stream) == 0);
check_cuda_error(cudaPeekAtLastError());

if constexpr (0) {
std::vector<uint32_t> tmp(input_dims * output_dims / 8);
Expand All @@ -194,8 +237,6 @@ struct Linear::Impl {
std::cout << std::hex << *it << " ";
}
}

weight_ = weight;
}

void convert_scales_zeros(void* workspace,
Expand All @@ -204,7 +245,8 @@ struct Linear::Impl {
size_t input_dims,
size_t output_dims,
int group_size,
bool use_simt)
bool use_simt,
cudaStream_t stream)
{
if constexpr (0) {
std::cout << "scales: " << std::endl;
Expand Down Expand Up @@ -245,12 +287,10 @@ struct Linear::Impl {
using namespace gemm;
auto [order_b, pack_b, order_v, pack_v] = get_weight_and_scales_layout(getSMVersion(), use_simt);

fuse_scales_and_zeros((half*)workspace, (const half*)scales.data, (half*)qzeros.data, scale_count);
sync_check_cuda_error();

cudaDeviceSynchronize();
fuse_scales_and_zeros((half*)workspace, (const half*)scales.data, (half*)qzeros.data, scale_count, stream);
check_cuda_error(cudaPeekAtLastError());

check_cuda_error(cudaMalloc(&scales_zeros_, sizeof(uint16_t) * scale_count * 2));
check_cuda_error(cudaMallocAsync(&scales_zeros_, sizeof(uint16_t) * scale_count * 2, stream));

MatrixLayout s_desc{
gemm::DataType::U32,
Expand All @@ -263,8 +303,8 @@ struct Linear::Impl {
q_desc_ = s_desc;
q_desc_.pack = pack_v;

TM_CHECK(Convert(workspace, s_desc, scales_zeros_, q_desc_, 0) == 0);
sync_check_cuda_error();
TM_CHECK(Convert(workspace, s_desc, scales_zeros_, q_desc_, stream) == 0);
check_cuda_error(cudaPeekAtLastError());

if constexpr (0) {
std::vector<__half> tmp(scale_count * 2);
Expand Down Expand Up @@ -300,14 +340,28 @@ struct Linear::Impl {
}

// create a new workspace if cache missed
auto workspace = std::shared_ptr<gemm::Workspace>(new gemm::Workspace, [device_id](gemm::Workspace* p) {
int old{};
check_cuda_error(cudaGetDevice(&old));
check_cuda_error(cudaSetDevice(device_id));
check_cuda_error(cudaFree(p->barriers));
check_cuda_error(cudaFree(p->partials));
check_cuda_error(cudaSetDevice(old));
});
auto workspace =
std::shared_ptr<gemm::Workspace>(new gemm::Workspace{}, [device_id](gemm::Workspace* p) noexcept {
int old_device{};
if (cudaGetDevice(&old_device) != cudaSuccess) {
delete p;
return;
}
if (old_device != device_id && cudaSetDevice(device_id) != cudaSuccess) {
delete p;
return;
}
// The buffers use cudaMallocAsync. Workspace teardown is outside
// the inference path, so synchronize before ordinary cudaFree.
if (cudaDeviceSynchronize() == cudaSuccess) {
cudaFree(p->barriers);
cudaFree(p->partials);
}
if (old_device != device_id) {
cudaSetDevice(old_device);
}
delete p;
});

workspace->barriers_size = gemm::Gemm::kBarriersSize;
workspace->partials_size = gemm::Gemm::kPartialsSize;
Expand All @@ -332,9 +386,10 @@ struct Linear::Impl {
size_t output_dims_;
int w_bit_;
int group_size_;
int device_id_{-1};

std::shared_ptr<Tensor> weight_;
half* scales_zeros_;
half* scales_zeros_{};

gemm::MatrixLayout k_desc_;
gemm::MatrixLayout q_desc_;
Expand All @@ -349,9 +404,10 @@ Linear::Linear(size_t input_dims, size_t output_dims, int w_bit, int group_size)
impl_ = std::make_shared<Impl>(input_dims, output_dims, w_bit, group_size);
}

void Linear::post_init(std::shared_ptr<Tensor> qweight, const Tensor& scales, const Tensor& qzeros, bool simt)
void Linear::post_init(
std::shared_ptr<Tensor> qweight, const Tensor& scales, const Tensor& qzeros, bool simt, cudaStream_t stream)
{
impl_->post_init(qweight, scales, qzeros, simt);
impl_->post_init(qweight, scales, qzeros, simt, stream);
}

void Linear::forward(const Tensor& in, Tensor& out, cudaStream_t stream)
Expand Down
6 changes: 5 additions & 1 deletion src/turbomind/api/python/linear.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,11 @@ enum class WeightType : int
class Linear {
public:
Linear(size_t input_dims, size_t output_dims, int w_bit, int group_size);
void post_init(std::shared_ptr<Tensor> qweight, const Tensor& scales, const Tensor& qzeros, bool simt);
void post_init(std::shared_ptr<Tensor> qweight,
const Tensor& scales,
const Tensor& qzeros,
bool simt,
cudaStream_t stream = nullptr);
void forward(const Tensor& in, Tensor& out, cudaStream_t stream = nullptr);
~Linear() {}

Expand Down
Loading