@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 122 lines 4.51 KB
Raw
1
2 #include "marlin.cuh"
3
4 #include <torch/csrc/stable/accelerator.h>
5 #include <torch/csrc/stable/library.h>
6 #include <torch/csrc/stable/ops.h>
7 #include <torch/csrc/stable/tensor.h>
8 #include <torch/headeronly/core/ScalarType.h>
9 #include <torch/headeronly/util/Exception.h>
10
11 #include "libtorch_stable/torch_utils.h"
12
13 // for only non-zp format (like gptq)
14 __global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
15 // qweight: (size_k * size_n // 8,)
16 const int32_t* __restrict__ qweight,
17 // output: same shape with qweight
18 int32_t* __restrict__ output) {
19 int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
20 int32_t new_val = 0;
21
22 #pragma unroll
23 for (int32_t i = 0; i < 8; i++) {
24 int32_t single_val = val & 0xF;
25 single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
26 new_val |= single_val << (i * 4);
27 val >>= 4;
28 }
29
30 output[blockIdx.x * 32 + threadIdx.x] = new_val;
31 }
32
33 // for awq format only (with zp and with awq weight layout)
34 __global__ void marlin_int4_fp8_preprocess_kernel_awq(
35 // AWQ qweight: (size_k, size_n // 8)
36 const int32_t* __restrict__ qweight,
37 // output: same shape with qweight
38 int32_t* __restrict__ output,
39 // AWQ zeros: (size_k // group_size, size_n // 8)
40 const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
41 int32_t group_size) {
42 // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
43 // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
44 // same row (stride 1) instead of striding across rows (stride size_n/8).
45 int col = blockIdx.y * 32 + threadIdx.x;
46 if (col >= size_n / 8) return;
47 (void)size_k;
48
49 int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
50 int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
51 int32_t new_val = 0;
52
53 #pragma unroll
54 for (int32_t i = 0; i < 8; i++) {
55 int32_t single_val = val & 0xF;
56 int32_t single_zero = zero & 0xF;
57
58 single_val =
59 single_val >= single_zero ? single_val - single_zero : 15 - single_val;
60 new_val |= single_val << (i * 4);
61 val >>= 4;
62 zero >>= 4;
63 }
64
65 output[blockIdx.x * (size_n / 8) + col] = new_val;
66 }
67
68 torch::stable::Tensor marlin_int4_fp8_preprocess(
69 torch::stable::Tensor& qweight,
70 std::optional<torch::stable::Tensor> qzeros_or_none, bool inplace) {
71 STD_TORCH_CHECK(qweight.is_cuda(), "qweight is not on GPU");
72 STD_TORCH_CHECK(qweight.scalar_type() == torch::headeronly::ScalarType::Int,
73 "qweight.dtype != torch.int32");
74
75 const int32_t device_index = qweight.get_device_index();
76 torch::stable::accelerator::DeviceGuard device_guard(device_index);
77 const cudaStream_t stream = get_current_cuda_stream(device_index);
78
79 torch::stable::Tensor output =
80 inplace ? qweight : torch::stable::empty_like(qweight);
81
82 if (!qzeros_or_none.has_value()) {
83 STD_TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
84 "qweight.numel() * 8 % 256 != 0");
85
86 int blocks = qweight.numel() * 8 / 256;
87 marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32, 0, stream>>>(
88 reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
89 reinterpret_cast<int32_t*>(output.mutable_data_ptr()));
90 } else {
91 int32_t size_k = qweight.size(0);
92 int32_t size_n = qweight.size(1) * 8;
93 torch::stable::Tensor qzeros = qzeros_or_none.value();
94
95 STD_TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
96 STD_TORCH_CHECK(qzeros.is_cuda(), "qzeros is not on GPU");
97 STD_TORCH_CHECK(qzeros.scalar_type() == torch::headeronly::ScalarType::Int,
98 "qweight.dtype != torch.int32");
99 STD_TORCH_CHECK(qzeros.get_device_index() == device_index,
100 "qzeros is not on the same device with qweight");
101
102 int32_t group_size = qweight.size(0) / qzeros.size(0);
103 STD_TORCH_CHECK(qweight.size(1) == qzeros.size(1),
104 "qweight.size(1) != qzeros.size(1)");
105 STD_TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
106 "qweight.size(0) % qzeros.size(0) != 0");
107 STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
108
109 dim3 blocks(size_k, (size_n / 8 + 31) / 32);
110 marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
111 reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
112 reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
113 reinterpret_cast<const int32_t*>(qzeros.const_data_ptr()), size_n,
114 size_k, group_size);
115 }
116
117 return output;
118 }
119
120 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
121 m.impl("marlin_int4_fp8_preprocess", TORCH_BOX(&marlin_int4_fp8_preprocess));
122 }