| 1 | #include "w4a8_utils.cuh" |
| 2 | |
| 3 | #include <array> |
| 4 | #include <cuda_runtime.h> |
| 5 | #include <cstdio> |
| 6 | |
| 7 | namespace vllm::cutlass_w4a8_utils { |
| 8 | |
| 9 | /* |
| 10 | GPU-accelerated implementation of cutlass::unified_encode_int4b. |
| 11 | Constructs a lookup table in constant memory to map 8 bits |
| 12 | (two 4-bit values) at a time. Assumes memory is contiguous |
| 13 | and pointers are 16-byte aligned. |
| 14 | */ |
| 15 | __constant__ uint8_t kNibbleLUT[256]; |
| 16 | |
| 17 | __global__ void unified_encode_int4b_device(const uint8_t* in, uint8_t* out, |
| 18 | size_t nbytes) { |
| 19 | constexpr size_t V = sizeof(uint4); // 16 bytes |
| 20 | const size_t tid = blockIdx.x * blockDim.x + threadIdx.x; |
| 21 | const size_t nthreads = size_t(gridDim.x) * blockDim.x; |
| 22 | const size_t nvec = nbytes / V; |
| 23 | |
| 24 | // 1-D grid-stride loop over 16-byte chunks |
| 25 | for (size_t vec = tid; vec < nvec; vec += nthreads) { |
| 26 | uint4 v = reinterpret_cast<const uint4*>(in)[vec]; |
| 27 | uint8_t* b = reinterpret_cast<uint8_t*>(&v); |
| 28 | #pragma unroll |
| 29 | for (int i = 0; i < int(V); ++i) b[i] = kNibbleLUT[b[i]]; |
| 30 | reinterpret_cast<uint4*>(out)[vec] = v; |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | static bool upload_lut() { |
| 35 | std::array<uint8_t, 256> lut{}; |
| 36 | auto map_nib = [](uint8_t v) -> uint8_t { |
| 37 | // 1..7 -> (8 - v); keep 0 and 8..15 |
| 38 | return (v == 0 || (v & 0x8)) ? v : uint8_t(8 - v); |
| 39 | }; |
| 40 | for (int b = 0; b < 256; ++b) { |
| 41 | uint8_t lo = b & 0xF; |
| 42 | uint8_t hi = (b >> 4) & 0xF; |
| 43 | lut[b] = uint8_t((map_nib(hi) << 4) | map_nib(lo)); |
| 44 | } |
| 45 | cudaError_t e = cudaMemcpyToSymbol(kNibbleLUT, lut.data(), lut.size(), |
| 46 | /*offset=*/0, cudaMemcpyHostToDevice); |
| 47 | |
| 48 | return (e == cudaSuccess); |
| 49 | } |
| 50 | |
| 51 | bool unified_encode_int4b(cutlass::int4b_t const* in, cutlass::int4b_t* out, |
| 52 | size_t num_int4_elems) { |
| 53 | // Build/upload LUT |
| 54 | if (!upload_lut()) return false; |
| 55 | |
| 56 | static_assert(sizeof(typename cutlass::int4b_t::Storage) == 1, |
| 57 | "int4 storage must be 1 byte"); |
| 58 | const size_t nbytes = num_int4_elems >> 1; |
| 59 | |
| 60 | auto* in_bytes = reinterpret_cast<uint8_t const*>(in); |
| 61 | auto* out_bytes = reinterpret_cast<uint8_t*>(out); |
| 62 | |
| 63 | // kernel launch params |
| 64 | constexpr int block = 256; |
| 65 | const size_t nvec = nbytes / sizeof(uint4); // # of 16B vectors |
| 66 | int grid = int((nvec + block - 1) / block); |
| 67 | if (grid == 0) grid = 1; // ensure we still cover the tail in the kernel |
| 68 | |
| 69 | unified_encode_int4b_device<<<grid, block>>>(in_bytes, out_bytes, nbytes); |
| 70 | |
| 71 | // launch errors |
| 72 | cudaError_t err = cudaGetLastError(); |
| 73 | if (err != cudaSuccess) { |
| 74 | printf("unified_encode_int4b_device launch error: %s (%d)\n", |
| 75 | cudaGetErrorString(err), err); |
| 76 | return false; |
| 77 | } |
| 78 | |
| 79 | // runtime errors |
| 80 | err = cudaDeviceSynchronize(); |
| 81 | if (err != cudaSuccess) { |
| 82 | printf("unified_encode_int4b_device runtime error: %s (%d)\n", |
| 83 | cudaGetErrorString(err), err); |
| 84 | return false; |
| 85 | } |
| 86 | |
| 87 | return true; |
| 88 | } |
| 89 | |
| 90 | } // namespace vllm::cutlass_w4a8_utils |