diff --git a/src/turbomind/api/python/bind.cpp b/src/turbomind/api/python/bind.cpp index db7ea2d..c716606 100644 --- a/src/turbomind/api/python/bind.cpp +++ b/src/turbomind/api/python/bind.cpp @@ -201,15 +201,20 @@ std::shared_ptr DLManagedTensorToTurbomindTensor(DLManagedTen return std::make_shared(where, dtype, shape, data); } -std::shared_ptr TorchTensorToTurbomindTensor(py::object obj) +std::shared_ptr 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(PyCapsule_GetPointer(cap.ptr(), kDlTensorCapsuleName)); return DLManagedTensorToTurbomindTensor(dlmt); } PYBIND11_MODULE(_turbomind_ext, m) { + m.attr("W4A16_LINEAR_API_VERSION") = 2; + py::enum_(m, "WeightType") .value("kFP32", turbomind::WeightType::kFP32) .value("kFP16", turbomind::WeightType::kFP16) @@ -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(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(stream_id); return self->forward(*_in, *_out, stream); }) diff --git a/src/turbomind/api/python/linear.cc b/src/turbomind/api/python/linear.cc index 0454cea..3b81de3 100644 --- a/src/turbomind/api/python/linear.cc +++ b/src/turbomind/api/python/linear.cc @@ -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 qweight, const Tensor& scales, const Tensor& qzeros, bool simt) + void post_init( + std::shared_ptr 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_); @@ -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, @@ -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 weight, size_t input_dims, size_t output_dims, bool use_simt) + void convert_qweight(void* workspace, + std::shared_ptr 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(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(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 tmp(input_dims * output_dims); @@ -173,12 +218,10 @@ struct Linear::Impl { k_desc_.type = gemm::DataType::U4; k_desc_.pack = pack_b; - cudaMemset(const_cast(weight->data), 0, input_dims * output_dims / 2); + check_cuda_error(cudaMemsetAsync(const_cast(weight->data), 0, input_dims * output_dims / 2, stream)); - TM_CHECK(Convert(workspace, w_desc, const_cast(weight->data), k_desc_, 0) == 0); - sync_check_cuda_error(); - - cudaDeviceSynchronize(); + TM_CHECK(Convert(workspace, w_desc, const_cast(weight->data), k_desc_, stream) == 0); + check_cuda_error(cudaPeekAtLastError()); if constexpr (0) { std::vector tmp(input_dims * output_dims / 8); @@ -194,8 +237,6 @@ struct Linear::Impl { std::cout << std::hex << *it << " "; } } - - weight_ = weight; } void convert_scales_zeros(void* workspace, @@ -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; @@ -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, @@ -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); @@ -300,14 +340,28 @@ struct Linear::Impl { } // create a new workspace if cache missed - auto workspace = std::shared_ptr(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(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; @@ -332,9 +386,10 @@ struct Linear::Impl { size_t output_dims_; int w_bit_; int group_size_; + int device_id_{-1}; std::shared_ptr weight_; - half* scales_zeros_; + half* scales_zeros_{}; gemm::MatrixLayout k_desc_; gemm::MatrixLayout q_desc_; @@ -349,9 +404,10 @@ Linear::Linear(size_t input_dims, size_t output_dims, int w_bit, int group_size) impl_ = std::make_shared(input_dims, output_dims, w_bit, group_size); } -void Linear::post_init(std::shared_ptr qweight, const Tensor& scales, const Tensor& qzeros, bool simt) +void Linear::post_init( + std::shared_ptr 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) diff --git a/src/turbomind/api/python/linear.h b/src/turbomind/api/python/linear.h index 1ebb3a7..0a494a5 100644 --- a/src/turbomind/api/python/linear.h +++ b/src/turbomind/api/python/linear.h @@ -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 qweight, const Tensor& scales, const Tensor& qzeros, bool simt); + void post_init(std::shared_ptr 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() {} diff --git a/tests/test_linear.py b/tests/test_linear.py new file mode 100644 index 0000000..979eda2 --- /dev/null +++ b/tests/test_linear.py @@ -0,0 +1,122 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import pytest +import torch + +import turbomind as tm +from turbomind import linear as linear_module +from turbomind.utils import get_u4_slices, pack_u4_row, unpack_awq_gemm + + +def _pack_nibbles(x: torch.Tensor) -> torch.Tensor: + """Pack the last dimension into int32 values, low nibble first.""" + assert x.dtype == torch.uint8 and x.shape[-1] == 8 + packed = torch.zeros(x.shape[:-1], dtype=torch.int64, device=x.device) + for index in range(8): + packed |= x[..., index].to(torch.int64) << (index * 4) + return packed.to(torch.int32) + + +def test_w4a16_linear_api_version_is_exported(): + assert 'W4A16_LINEAR_API_VERSION' in tm.__all__ + assert isinstance(tm.W4A16_LINEAR_API_VERSION, int) + + if linear_module.TURBOMIND_EXTENSION_INSTALLED: + assert tm.W4A16_LINEAR_API_VERSION == 2 + assert tm.W4A16_LINEAR_API_VERSION == getattr( + linear_module._turbomind_ext, 'W4A16_LINEAR_API_VERSION') + else: + assert tm.W4A16_LINEAR_API_VERSION == 0 + + +def test_linear_reports_missing_native_extension(monkeypatch): + monkeypatch.setattr(linear_module, 'TURBOMIND_EXTENSION_INSTALLED', False) + + with pytest.raises(RuntimeError, match='_turbomind_ext'): + tm.Linear(128, + 32, + bias=False, + quant_method='awq', + w_bit=4, + group_size=128, + device='cpu') + + +def test_awq_unpack_and_private_row_pack_layout(): + storage_nibbles = torch.arange(8, dtype=torch.uint8).view(1, 1, 8) + awq_packed = _pack_nibbles(storage_nibbles) + + logical = unpack_awq_gemm(awq_packed) + expected_logical = storage_nibbles[..., + [0, 4, 1, 5, 2, 6, 3, 7]].reshape(1, 8) + torch.testing.assert_close(logical, expected_logical) + + private_packed = pack_u4_row(logical) + private_nibbles = torch.stack(get_u4_slices(private_packed, torch.uint8), + dim=-1) + torch.testing.assert_close(private_nibbles, expected_logical.view(1, 1, 8)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA is required') +def test_native_awq_linear_cross_stream(): + if not linear_module.TURBOMIND_EXTENSION_INSTALLED: + pytest.skip('_turbomind_ext is not built') + if tm.W4A16_LINEAR_API_VERSION < 2: + pytest.fail('Rebuild _turbomind_ext with W4A16 Linear API v2') + + device = torch.device('cuda:0') + in_features = 128 + out_features = 128 + group_size = 128 + + generator = torch.Generator(device=device).manual_seed(0) + weight_nibbles = torch.randint(0, + 16, (in_features, out_features // 8, 8), + dtype=torch.uint8, + device=device, + generator=generator) + zero_nibbles = torch.randint( + 0, + 16, (in_features // group_size, out_features // 8, 8), + dtype=torch.uint8, + device=device, + generator=generator) + qweight = _pack_nibbles(weight_nibbles) + qzeros = _pack_nibbles(zero_nibbles) + scales = torch.rand((in_features // group_size, out_features), + dtype=torch.float16, + device=device, + generator=generator) * 0.1 + + unpacked_weight = unpack_awq_gemm(qweight).float() + unpacked_zeros = unpack_awq_gemm(qzeros).float() + reference_weight = ((unpacked_weight - unpacked_zeros) * + scales.float()).half() + + op = tm.Linear(in_features, + out_features, + bias=False, + quant_method='awq', + w_bit=4, + group_size=group_size, + device=device) + op.qweight.copy_(qweight) + op.qzeros.copy_(qzeros) + op.scales.copy_(scales) + + init_stream = torch.cuda.Stream(device=device) + run_stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(init_stream): + op.post_init() + + x = torch.randn((5, in_features), + dtype=torch.float16, + device=device, + generator=generator) + with torch.cuda.stream(run_stream): + actual = op(x) + expected = x @ reference_weight + run_stream.synchronize() + + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=2e-2) + tm.Linear.clear_workspaces() diff --git a/turbomind/__init__.py b/turbomind/__init__.py index 5de11b0..e226a3d 100644 --- a/turbomind/__init__.py +++ b/turbomind/__init__.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -from .linear import Linear +from .linear import W4A16_LINEAR_API_VERSION, Linear -__all__ = ['Linear'] +__all__ = ['Linear', 'W4A16_LINEAR_API_VERSION'] diff --git a/turbomind/linear.py b/turbomind/linear.py index 5d44317..721d7f8 100644 --- a/turbomind/linear.py +++ b/turbomind/linear.py @@ -15,9 +15,12 @@ try: import _turbomind_ext TURBOMIND_EXTENSION_INSTALLED = True + W4A16_LINEAR_API_VERSION = getattr(_turbomind_ext, + 'W4A16_LINEAR_API_VERSION', 0) except Exception as e: logging.error(f'_turbomind_ext is not installed: {e}') TURBOMIND_EXTENSION_INSTALLED = False + W4A16_LINEAR_API_VERSION = 0 def pad_out_dims(x: torch.Tensor, dims: int): @@ -61,6 +64,11 @@ def __init__(self, self.bias = bias self.quant_method = quant_method + if not TURBOMIND_EXTENSION_INSTALLED: + raise RuntimeError( + 'turbomind kernels are not installed. Build/install the ' + '_turbomind_ext extension before constructing Linear.') + # quick sanity check (make sure alignment) assert self.in_features % self.group_size == 0 assert self.out_features % (32 // self.w_bit) == 0 @@ -106,8 +114,14 @@ def __init__(self, self.linear = _turbomind_ext.Linear(self.in_features, self.out_features, self.w_bit, self.group_size) + self._prepared = False + self._weight_owner = None + self._ready_event = None + self._ready_streams = set() def post_init(self): + if self._prepared: + raise RuntimeError('turbomind.Linear.post_init() is one-shot.') assert self.qweight.device.type == 'cuda' if self.quant_method == 'awq': self.qweight = unpack_awq_gemm(self.qweight) @@ -124,7 +138,8 @@ def post_init(self): self.qzeros = self.qzeros self.scales = self.scales else: - return + raise ValueError( + f'Unsupported quant_method: {self.quant_method!r}') self.qweight = pack_u4_row(self.qweight) self.qzeros = self.qzeros.to(torch.half) @@ -141,10 +156,26 @@ def is_16xx_series(name): self.qweight = self.qweight.contiguous() self.scales = self.scales.contiguous() self.qzeros = self.qzeros.contiguous() - self.linear.post_init(self.qweight, self.scales, self.qzeros, simt) + with torch.cuda.device(self.qweight.device): + stream = torch.cuda.current_stream() + try: + self.linear.post_init(self.qweight, self.scales, self.qzeros, + simt, stream.cuda_stream) + finally: + self.qweight.record_stream(stream) + self.scales.record_stream(stream) + self.qzeros.record_stream(stream) + self._weight_owner = self.qweight + self._ready_event = torch.cuda.Event() + self._ready_event.record(stream) + self._ready_streams.add(stream.cuda_stream) + self._prepared = True @torch.no_grad() def forward(self, x): + if not self._prepared: + raise RuntimeError( + 'turbomind.Linear is not prepared; call post_init() first.') assert TURBOMIND_EXTENSION_INSTALLED, ( 'turbomind kernels are not installed. ' 'Please perform `pip install turbomind` to install turbomind ' @@ -159,9 +190,26 @@ def forward(self, x): dtype=torch.float16, device=x.device, ) - stream = torch.cuda.current_stream() - self.linear.forward(x, out, stream.cuda_stream) - out = torch.from_dlpack(out) + with torch.cuda.device(x.device): + stream = torch.cuda.current_stream() + stream_id = stream.cuda_stream + if self._ready_event is not None: + capturing = torch.cuda.is_current_stream_capturing() + if not capturing and self._ready_event.query(): + self._ready_event = None + self._ready_streams.clear() + elif stream_id not in self._ready_streams: + if capturing: + raise RuntimeError( + 'Run one eager forward on this CUDA stream before ' + 'capturing turbomind.Linear.') + stream.wait_event(self._ready_event) + self._ready_streams.add(stream_id) + try: + self.linear.forward(x, out, stream.cuda_stream) + finally: + x.record_stream(stream) + self._weight_owner.record_stream(stream) if self.bias is not None: out.add_(self.bias) return out.view(out_shape)