From 23b39c97bc328e7f997fd63e19f7792548e2f07a Mon Sep 17 00:00:00 2001 From: qescccczmr Date: Mon, 24 Aug 2026 03:41:28 +0000 Subject: [PATCH] feat(linear): expose stream-aware W4A16 API Propagate the caller CUDA stream through DLPack, weight preparation, and GEMM dispatch while managing asynchronous buffer lifetimes safely. Export an API capability version and cover layout conversion plus native cross-stream execution. --- src/turbomind/api/python/bind.cpp | 29 ++++-- src/turbomind/api/python/linear.cc | 138 ++++++++++++++++++++--------- src/turbomind/api/python/linear.h | 6 +- tests/test_linear.py | 122 +++++++++++++++++++++++++ turbomind/__init__.py | 4 +- turbomind/linear.py | 58 ++++++++++-- 6 files changed, 299 insertions(+), 58 deletions(-) create mode 100644 tests/test_linear.py 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)