Skip to content

Commit efa4f0f

Browse files
spike-zhuspike-zhu
authored
issue/453: add AWQ-dequantize in iluvatar gpu
1 parent fbfb0ef commit efa4f0f

5 files changed

Lines changed: 209 additions & 0 deletions

File tree

‎src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cu‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
#include "../../../devices/nvidia/nvidia_handle.cuh"
2+
#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
3+
#include "dequantize_w42f16_iluvatar.cuh"
4+
#include "dequantize_w42f16_kernel.cuh"
5+
6+
#include "../dequantize_awq.h"
7+
#include <cuda_fp16.h>
8+
9+
__global__ void __launch_bounds__(64)
10+
dequantize_weights(int *__restrict__ B, half *__restrict__ scaling_factors,
11+
int *__restrict__ zeros, half *__restrict__ C, int G) {
12+
// static constexpr uint32_t ZERO = 0x0;
13+
half B_shared[32 * (128 + 8)];
14+
15+
half *B_shared_ptr2 = B_shared;
16+
17+
int N = blockDim.x * gridDim.x; // 2
18+
int col = (blockIdx.x * blockDim.x + threadIdx.x);
19+
int row = (blockIdx.y * blockDim.y + threadIdx.y);
20+
int index1 = 8 * col + 8 * row * N;
21+
half *C_ptr2 = C + index1;
22+
23+
int index2 = col + row * N;
24+
int *B_ptr2 = B + index2;
25+
26+
int index3 = col + (int)(row / G) * N;
27+
int *zeros_ptr2 = zeros + index3;
28+
int index4 = 8 * col + (int)(row / G) * N * 8;
29+
half *scaling_factors_ptr2 = scaling_factors + index4;
30+
31+
uint32_t zeros_loaded = *(uint32_t *)(zeros_ptr2);
32+
uint4 B_loaded_zero = dequantize_s4_to_fp16x2(zeros_loaded);
33+
uint4 B_loaded_scale = *(uint4 *)(scaling_factors_ptr2);
34+
35+
uint32_t B_loaded = *(uint32_t *)B_ptr2;
36+
uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2(B_loaded);
37+
38+
// Reinterpret uint4 components as __half2
39+
__half2 *B_loaded_fp16_h2 = reinterpret_cast<__half2 *>(&B_loaded_fp16);
40+
__half2 *B_loaded_zero_h2 = reinterpret_cast<__half2 *>(&B_loaded_zero);
41+
__half2 *B_loaded_scale_h2 = reinterpret_cast<__half2 *>(&B_loaded_scale);
42+
43+
// Replace PTX sub.f16x2 with __hsub2 for each component
44+
B_loaded_fp16_h2[0] = __hsub2(B_loaded_fp16_h2[0], B_loaded_zero_h2[0]);
45+
B_loaded_fp16_h2[1] = __hsub2(B_loaded_fp16_h2[1], B_loaded_zero_h2[1]);
46+
B_loaded_fp16_h2[2] = __hsub2(B_loaded_fp16_h2[2], B_loaded_zero_h2[2]);
47+
B_loaded_fp16_h2[3] = __hsub2(B_loaded_fp16_h2[3], B_loaded_zero_h2[3]);
48+
49+
// Replace PTX fma.rn.f16x2 with __hfma2 for each component
50+
B_loaded_fp16_h2[0] = __hfma2(B_loaded_fp16_h2[0], B_loaded_scale_h2[0], __float2half2_rn(0.0f));
51+
B_loaded_fp16_h2[1] = __hfma2(B_loaded_fp16_h2[1], B_loaded_scale_h2[1], __float2half2_rn(0.0f));
52+
B_loaded_fp16_h2[2] = __hfma2(B_loaded_fp16_h2[2], B_loaded_scale_h2[2], __float2half2_rn(0.0f));
53+
B_loaded_fp16_h2[3] = __hfma2(B_loaded_fp16_h2[3], B_loaded_scale_h2[3], __float2half2_rn(0.0f));
54+
55+
// Store back to shared memory
56+
*(uint4 *)B_shared_ptr2 = B_loaded_fp16;
57+
58+
for (int i = 0; i < 8; ++i) {
59+
*(C_ptr2 + i) = B_shared[i];
60+
}
61+
}
62+
63+
namespace op::dequantize_awq::iluvatar {
64+
65+
struct Descriptor::Opaque {
66+
std::shared_ptr<device::nvidia::Handle::Internal> internal;
67+
};
68+
69+
Descriptor::~Descriptor() {
70+
delete _opaque;
71+
}
72+
73+
infiniStatus_t Descriptor::create(
74+
infiniopHandle_t handle_,
75+
Descriptor **desc_ptr,
76+
infiniopTensorDescriptor_t out_desc,
77+
infiniopTensorDescriptor_t qweight_desc,
78+
infiniopTensorDescriptor_t scales_desc,
79+
infiniopTensorDescriptor_t zeros_desc) {
80+
81+
auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
82+
auto result = DequantizeAWQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc);
83+
84+
*desc_ptr = new Descriptor(
85+
0,
86+
new Opaque{handle->internal()},
87+
result.take(),
88+
handle->device, handle->device_id);
89+
return INFINI_STATUS_SUCCESS;
90+
}
91+
92+
infiniStatus_t
93+
Descriptor::calculate(
94+
void *workspace,
95+
size_t workspace_size,
96+
void *out,
97+
const void *qweight,
98+
const void *scales,
99+
const void *zeros,
100+
void *stream) const {
101+
int in_features = _info.in_features();
102+
int out_features = _info.out_features();
103+
int group_size = in_features / _info.num_groups();
104+
105+
// ==================== 默认配置, 固定为 8 ====================
106+
constexpr int BLOCK_X = 8;
107+
constexpr int BLOCK_Y = 8;
108+
109+
int x_blocks = (out_features + BLOCK_X - 1) / BLOCK_X;
110+
int y_blocks = (in_features + BLOCK_Y - 1) / BLOCK_Y;
111+
112+
dim3 num_blocks(x_blocks, y_blocks);
113+
dim3 threads_per_block(BLOCK_X, BLOCK_Y);
114+
// =====================================================
115+
116+
half *out_ = reinterpret_cast<half *>(out);
117+
118+
int *qweight_ = const_cast<int *>(reinterpret_cast<const int *>(qweight));
119+
half *scales_ = const_cast<half *>(reinterpret_cast<const half *>(scales));
120+
int *zeros_ = const_cast<int *>(reinterpret_cast<const int *>(zeros));
121+
122+
dequantize_weights<<<num_blocks, threads_per_block, 0, reinterpret_cast<cudaStream_t>(stream)>>>(
123+
qweight_, scales_, zeros_, out_, group_size);
124+
return INFINI_STATUS_SUCCESS;
125+
}
126+
127+
} // namespace op::dequantize_awq::iluvatar

‎src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cuh‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#ifndef __DEQUANTIZE_AWQ_ILUVATAR_CUH__
2+
#define __DEQUANTIZE_AWQ_ILUVATAR_CUH__
3+
4+
#include "../dequantize_awq.h"
5+
6+
DESCRIPTOR(iluvatar)
7+
8+
#endif // __DEQUANTIZE_AWQ_ILUVATAR_CUH__

‎src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_kernel.cuh‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#pragma once
2+
#include <cuda_fp16.h> // 需要此头文件来支持 __half 和 __half2 类型
3+
4+
/**
5+
* @brief 将一个包含8个4-bit整数的uint32_t反量化为8个half精度浮点数。
6+
*
7+
* 这是一个通用的 CUDA C++ 实现，用于替代原有的 PTX 汇编版本，
8+
* 以便在不支持高级 PTX 指令（如 lop3.b32）的 GPU 上运行。
9+
* 输出顺序匹配 PTX 的交错打包：v0, v4, v1, v5, v2, v6, v3, v7（经 signed 调整后）。
10+
*
11+
* @param source 输入的32位无符号整数，它打包了8个4-bit的数据。
12+
* @return 一个 uint4 变量，其中包含8个反量化后的 half 值。
13+
*/
14+
__device__ __forceinline__ uint4 dequantize_s4_to_fp16x2(uint32_t const &source) {
15+
// 步骤 1: 从一个 32-bit 源数据中解包出 8 个 4-bit 无符号整数。
16+
// 源数据的内存布局被假定为 [v7, v6, v5, v4, v3, v2, v1, v0]，
17+
// 其中每个 'v' 都是一个 4-bit 的半字节 (nibble)。
18+
const unsigned int v0 = (source >> 0) & 0x0F;
19+
const unsigned int v1 = (source >> 4) & 0x0F;
20+
const unsigned int v2 = (source >> 8) & 0x0F;
21+
const unsigned int v3 = (source >> 12) & 0x0F;
22+
const unsigned int v4 = (source >> 16) & 0x0F;
23+
const unsigned int v5 = (source >> 20) & 0x0F;
24+
const unsigned int v6 = (source >> 24) & 0x0F;
25+
const unsigned int v7 = (source >> 28) & 0x0F;
26+
27+
// 步骤 2: 对于 signed 4-bit (s4)，减去 8 以映射到 [-8, 7] 范围。
28+
// 定义偏移量
29+
__half offset = __half(8);
30+
31+
// 计算 signed 值
32+
__half hv0 = __half(v0) - offset;
33+
__half hv1 = __half(v1) - offset;
34+
__half hv2 = __half(v2) - offset;
35+
__half hv3 = __half(v3) - offset;
36+
__half hv4 = __half(v4) - offset;
37+
__half hv5 = __half(v5) - offset;
38+
__half hv6 = __half(v6) - offset;
39+
__half hv7 = __half(v7) - offset;
40+
41+
// 步骤 3: 将 half 值按 PTX 交错顺序打包成 __half2 并存入 result 中。
42+
// 顺序：result_ptr[0]: low=hv0, high=hv4
43+
// result_ptr[1]: low=hv1, high=hv5
44+
// result_ptr[2]: low=hv2, high=hv6
45+
// result_ptr[3]: low=hv3, high=hv7
46+
// __halves2half2 函数：low 为第一个参数，high 为第二个参数。
47+
uint4 result;
48+
__half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
49+
50+
result_ptr[0] = __halves2half2(hv0, hv4);
51+
result_ptr[1] = __halves2half2(hv1, hv5);
52+
result_ptr[2] = __halves2half2(hv2, hv6);
53+
result_ptr[3] = __halves2half2(hv3, hv7);
54+
55+
return result;
56+
}

‎src/infiniop/ops/dequantize_awq/operator.cc‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
#ifdef ENABLE_MOORE_API
99
#include "moore/dequantize_w42f16_moore.h"
1010
#endif
11+
#ifdef ENABLE_ILUVATAR_API
12+
#include "iluvatar/dequantize_w42f16_iluvatar.cuh"
13+
#endif
1114

1215
__C infiniStatus_t infiniopCreateDequantizeAWQDescriptor(
1316
infiniopHandle_t handle,
@@ -33,6 +36,9 @@ __C infiniStatus_t infiniopCreateDequantizeAWQDescriptor(
3336
#endif
3437
#ifdef ENABLE_MOORE_API
3538
CREATE(INFINI_DEVICE_MOORE, moore);
39+
#endif
40+
#ifdef ENABLE_ILUVATAR_API
41+
CREATE(INFINI_DEVICE_ILUVATAR, iluvatar);
3642
#endif
3743
default:
3844
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -54,6 +60,9 @@ __C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDe
5460
#endif
5561
#ifdef ENABLE_MOORE_API
5662
GET(INFINI_DEVICE_MOORE, moore);
63+
#endif
64+
#ifdef ENABLE_ILUVATAR_API
65+
GET(INFINI_DEVICE_ILUVATAR, iluvatar);
5766
#endif
5867
default:
5968
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -82,6 +91,9 @@ __C infiniStatus_t infiniopDequantizeAWQ(
8291
#endif
8392
#ifdef ENABLE_MOORE_API
8493
CALCULATE(INFINI_DEVICE_MOORE, moore);
94+
#endif
95+
#ifdef ENABLE_ILUVATAR_API
96+
CALCULATE(INFINI_DEVICE_ILUVATAR, iluvatar);
8597
#endif
8698
default:
8799
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -104,6 +116,9 @@ infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc) {
104116
#endif
105117
#ifdef ENABLE_MOORE_API
106118
DELETE(INFINI_DEVICE_MOORE, moore);
119+
#endif
120+
#ifdef ENABLE_ILUVATAR_API
121+
DELETE(INFINI_DEVICE_ILUVATAR, iluvatar);
107122
#endif
108123
default:
109124
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

‎xmake/iluvatar.lua‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,9 @@ target("infiniop-iluvatar")
5050
-- set_languages("cxx17") 天数似乎不能用这个配置
5151
add_files("../src/infiniop/devices/nvidia/*.cu", "../src/infiniop/ops/*/nvidia/*.cu")
5252

53+
-- 天数平台不支持部分 NVIDIA PTX 指令，AWQ 反量化改用 CUDA C++ 实现
54+
add_files("../src/infiniop/ops/dequantize_awq/iluvatar/*.cu")
55+
5356
if has_config("ninetoothed") then
5457
add_files("../build/ninetoothed/*.c", {cxflags = {"-Wno-return-type"}})
5558
end

0 commit comments

Comments
 (0)