DeepSeek 本地部署——蒸馏版、量化版和满血版实测效果对比

引言

随着 DeepSeek 受到世界范围广泛关注，官网服务经常不可用，除了调用第三方提供的 API 之外，本地部署也成为替代方案之一。

本文将介绍本地部署 DeepSeek 的三种不同类型模型：满血版，1.58 bit 量化版和蒸馏版（Qwen 1.5B），使用同一代码生成问题进行效果评测，最后给出实际部署建议。

硬件需求

以上只是推荐配置，随着社区和硬件厂商不断优化，将来可能会有更具性价比的部署方案。

部署过程

部署过程较冗长，本文暂不展开介绍，有需求可以加微信：Mistake113。

效果实测

让 DeepSeek 写一段 CUDA 代码，实现 Transformer 架构中必不可少的 Softmax 算子。

先看蒸馏版：

<思考过程>

嗯，用户让我写一段CUDA代码来实现softmax功能。首先，我得回忆一下什么是Softmax函数。它通常用于在多分类的问题中确保概率分布的总和为1，并且每个元素都大于0。

然后，考虑CUDA的特点。CUDA是基于GPU的加速技术，适合处理并行计算。所以，我的代码必须设计成可以高效利用GPU的parallelism。

接下来，我需要分析Softmax的具体步骤。一般来说，Softmax会先将每个输入向量减去最大的值（归一化），然后指数化，最后除以这些指数的和。这一步是为了避免数值溢出，并且确保结果是概率分布。

在并行处理方面，我的代码可能需要将输入数据分割到不同的块，然后分别计算各自的归一化因子。每个block中的元素可以用 shared 和共享 Memory 来高效地进行通信，减少显存访问次数。

同时，我还需要考虑数据的格式。软体通常以矩阵的形式存储，所以每个CudaSharedMemoryBlock可能需要一个float matrix来存储归一化因子。然后计算所有块中的归一化因子之和，最后对结果进行归一化处理。

此外，可能会有一些边界条件需要注意，比如当所有的元素相同时的情况，或者非常接近的值导致数值不稳定。但在CUDA中可以使用 fused 可变参数来解决这些问题。

最后，我需要将这些步骤组织成一个高效的并行代码块，确保逻辑清晰，结构完整，并且能够处理各种输入情况。 

<结束思考>

#include <math.h>#include <stdio.h>#include <string>
// 共享内存 Block 类class CudaSharedMemoryBlock {private:std::vector<std::vector<float>> shared;public:CudaSharedMemoryBlock() : _shared{} {}
void push(std::vector<std::vector<float>>& out) {for (size_t i = 0; i < out.size(); ++i) {for (size_t j = 0; j < out[i].size(); ++j) {out[i][j] = _shared shared[i][j];}}}
size_t get_size() const { return _shared.size(); }};
// Softmax 函数 实现void softmax(float *input, float *output, size_t block_size, int num_blocks) {CudaSharedMemoryBlock smb_block;
// 为所有块分配共享内存for (size_t i = 0; i < num_blocks; ++i) {const float* out = new float[num_blocks];if (!sobsm_block.get_size() == output.size()) {*out = 0.0f;}// 在共享内存中存储归一化因子for (size_t j = 0; j < num_blocks; ++j) {out[j] = smb_block.push(sobsm_block.get_size());}
// 对共享内存中的值取指数for (size_t i = 0; i < num_blocks; ++i) {const float* sum = new float[num_blocks];if (!sobsm_block.get_size() == sum.size()) {*sum = 0.0f;}// 计算归一化因子for (size_t j = 0; j < num_blocks; ++j) {sum[j] = expf(sobsm_block.push(sobsm_block.get_size(), smb_block.get_size() * i));}
// 计算归一化因子之和float total_sum = 0.0f;for (size_t j = 0; j < num_blocks; ++j) {total_sum += sum[j];}
// 对归一化因子进行归一化处理if (total_sum == 0.0f)return;
for (size_t i = 0; i < num_blocks; ++i) {const float* normalized = new float[num_blocks];if (!sobsm_block.get_size() == normalized.size()) {*normalized = 0.0f;}for (size_t j = 0; j < num_blocks; ++j) {normalized[j] = smb_block.push(sobsm_block.get_size(), total_sum, i);normalized[j] /= total_sum;}
// 将结果释放for (size_t i = 0; i < num_blocks; ++i) {delete[] out[i];delete[] sum[i];delete[] normalized[i];}}

蒸馏版思考过程最短，思路也最简单，问题是只实现了 C++ 代码，没有写 CUDA Kernel。另外生成的代码编译时报错：

笔者懒得去手动修复，直接放弃。

之后我们看满血版（671B），解决这个代码问题足足思考了 2 分钟：

<思考过程>
嗯，用户让我写一段CUDA代码实现softmax功能。首先，我需要确认softmax的具体实现方式，以及如何在CUDA中高效地实现它。Softmax通常用于将一个向量中的数值转换为概率分布，每个元素的值在0到1之间，并且所有元素的和为1。那，softmax的计算步骤是对每个元素求指数，然后除以所有元素的指数和。公式是softmax(x_i) = exp(x_i) / sum(exp(x))。在CUDA中实现的话，需要考虑并行化和数值稳定性问题。数值稳定性方面，通常会对输入数据减去最大值，防止指数运算时数值溢出。所以需要先找出每个数据块中的最大值，然后每个元素减去这个最大值，再进行指数运算。这样处理之后，计算会更稳定。接下来，问题是如何在CUDA中高效地计算每个元素。因为softmax通常是在一个向量上操作，比如在神经网络的最后一层，假设输入是一个二维张量，形状是(batch_size, feature_size)，我们需要对每个batch中的feature进行softmax处理。所以在这个情况下，每个线程块处理一个batch样本，或者共享内存来加速计算。

但在实际情况中，可能需要处理更大的数据，所以需要考虑如何分块。比如，对于每个输入向量来说，计算分为两步：第一步是计算最大值，第二步是计算指数和，第三步是每个元素除以总和。

在CUDA中，通常使用归约（reduction）操作来计算最大值和总和。所以可能需要编写两个核函数：一个计算最大值，另一个计算总和。或者，可以在同一个核函数中分两步处理。另外，还可以利用现有的CUDA库，例如CUB库中的BlockReduce来简化归约操作。

不过用户可能希望自己手动实现，所以需要考虑如何分块处理。例如，对于一个向量，每个线程块处理一个向量，块内的线程协作计算最大值和总和。

现在，具体实现步骤可能如下：

1. 确定输入的维度，假设输入是一维数组，或者二维数组的batch处理。比如，输入是形状为(n, m)的矩阵，对每行进行softmax处理。这里假设每个线程块处理一行，每个线程处理该行中的一个元素。

2. 核函数的参数包括输入数组、输出数组，以及每行的元素数量。比如，输入数组d_x，输出数组d_y，每行有cols个元素。

3. 在核函数中，首先确定当前线程块处理的行号。比如，blockIdx.x对应行号，blockDim.x是线程块中的线程数，线程的索引是threadIdx.x。

4. 对于每行，首先寻找最大值。这需要所有线程协作进行归约。每个线程先读取自己负责的元素（如果存在），然后进行块内的最大值归约。例如，使用共享内存来存储中间结果，逐层归约得到最大值。

5. 得到最大值后，每个线程计算其元素的指数，同时累加这些指数值。这时候同样需要归约求和。这两个归约（最大值和求和）可以分两个阶段进行。

6. 最后，每个线程再次计算指数后的值，并除以总和，得到最终的softmax结果。

但是需要注意的是，每一步都需要同步线程，确保所有线程完成当前阶段的计算，特别是在共享内存操作时。

可能的数值问题：当处理非常大的数值时，指数运算可能导致溢出。所以减去最大值是必要的。例如，计算每个元素的x_i - max_x，然后求指数，这样可以保证最大的指数是exp(0)=1，其他都是更小的值，不会溢出。

所以具体步骤：

对于每个行向量：

1. 找到该行的最大值max_val。

2. 计算每个元素的指数，并累加得到sum_exp。

3. 每个元素指数除以sum_exp。

现在，如何在CUDA中实现这些步骤？

首先，对于每行，启动一个线程块。假设每行有cols元素，线程块中的线程数目设为cols，或者更实际的情况是，当cols较大时，如1024，但CUDA线程块最大线程数是1024，所以对于超过的情况，需要用多个线程循环处理。

假设每行的元素数量是cols，每个线程块处理一行，每个线程处理多个元素。或者，如果cols比较大，可以将线程块的大小设为256或512，并让每个线程处理多个元素，通过循环来处理。例如，每个线程处理cols / blockDim.x 个元素。

不过为了简化，假设每行的元素数量不超过1024，这样每个线程块可以处理一行，每个线程处理一个元素。这会更容易实现。

例如，核函数可能的结构：

global void softmax_kernel(float *d_x, float *d_y, int cols) {

extern __shared__ float shared_mem[]; // 共享内存用于归约 int row = blockIdx.x; int tid = threadIdx.x; // 计算该行的起始指针 float *x_row = d_x + row * cols; float *y_row = d_y + row * cols; // 第一步：找到最大值 float max_val = -INFINITY; for (int i = tid; i < cols; i += blockDim.x) { float val = x_row[i]; if (val > max_val) max_val = val; } // 将每个线程的max_val放入共享内存，进行块内归约 shared_mem[tid] = max_val; __syncthreads(); // 归约求最大值，比如使用树状归约 for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { if (shared_mem[tid + s] > shared_mem[tid]) { shared_mem[tid] = shared_mem[tid + s]; } } __syncthreads(); } max_val = shared_mem[0]; __syncthreads(); // 第二步：计算指数和 float exp_sum = 0.0f; for (int i = tid; i < cols; i += blockDim.x) { exp_sum += expf(x_row[i] - max_val); } // 将每个线程的exp_sum累加到共享内存 shared_mem[tid] = exp_sum; __syncthreads(); // 归约求和 for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { shared_mem[tid] += shared_mem[tid + s]; }

__syncthreads(); } exp_sum = shared_mem[0]; __syncthreads(); // 第三步：计算每个元素的值并写入结果 if (exp_sum == 0) { // 避免除以零，可能需要处理异常情况，这里假设exp_sum不为零 } else { for (int i = tid; i < cols; i += blockDim.x) { y_row[i] = expf(x_row[i] - max_val) / exp_sum; } }

}

}

这里有几个问题需要注意：

1. 共享内存的大小需要足够，这里假设每个线程块中的线程数量等于共享内存的大小。所以在启动核函数时，动态分配共享内存的大小，比如每个线程块需要blockDim.x * sizeof(float)的共享内存。在核函数调用时，第三个参数就是所需的共享内存大小。

例如，当启动该核函数时：

int threads_per_block = 256;

int shared_mem_size = threads_per_block * sizeof(float);

softmax_kernel<<<num_rows, threads_per_block, shared_mem_size>>>(d_x, d_y, cols);

1. 每行的元素数cols可能大于线程块中的线程数。例如，如果cols是1000，而threads_per_block是256，那么每个线程需要处理多个元素，使用循环，比如i += blockDim.x的方式。

2. 第一阶段的归约是正确的，但需要确保所有线程都参与归约。例如，当blockDim.x不是2的幂时，树状归约可能需要调整。或者使用更稳定的归约方法，例如使用CUB的BlockReduce。

不过，这里为了简化，假设threads_per_block是2的幂，比如256，512，或1024。这可能要求调用者根据cols的大小调整线程块配置。

另外，处理每行的线程块的数量应该是num_rows，即每个线程块处理一行。

例如，假设输入是一个二维数组，形状为(rows, cols)，则核函数启动配置是：

dim3 grid(rows);

dim3 block(threads_per_block);

这样，每个线程块处理一行，每个线程处理该行中的一个或多个元素。

此外，在计算指数和的时候，可能需要注意精度问题，使用expf而不是exp，因为输入数据是float类型。

另外，关于数值稳定性，代码中已经减去了max_val，所以没问题。

现在，考虑这个核函数的正确性。例如，对于每行，每个线程处理多个元素的情况是否会导致归约错误？

比如，在最大值计算阶段，每个线程的循环覆盖所有该行的元素，计算自己处理过的元素中的最大值，然后将这些最大值进行块内归约。这将得到该行的最大值。

在归约结束后，所有线程都得到相同的max_val，这用于后续的指数运算。

同样，对于指数和的计算，每个线程计算自己处理的元素的指数之和，然后归约得到总和，这样整个行的总和被正确计算。

最后，每个线程再次处理自己负责的元素，计算exp(x_i - max_val) / sum_exp，得到最终的softmax值。

这样的实现是正确的，但在效率上可能还有优化空间，例如使用更高效的归约方法，或者利用共享内存的bank冲突避免等。

另外，当cols很大的时候，例如超过了每个线程块的最大线程数，这时候线程块中的线程可能需要处理多个元素，这可能需要多个循环步骤。例如，每个线程处理cols / blockDim.x 个元素，向上取整。

例如：

int tid = threadIdx.x;

int num_elements_per_thread = (cols + blockDim.x - 1) / blockDim.x;

for (int idx = 0; idx < num_elements_per_thread; idx++) {

int i = tid + idx * blockDim.x; if (i < cols) { // 处理第i个元素 }

}

这可能更高效，特别是当cols远大于blockDim.x时。

所以，可能需要将核函数中的循环部分进行相应的调整。比如，把每个线程处理的元素数量动态计算。

例如，在最大值计算的循环中：

for (int idx = 0; idx < num_elements_per_thread; idx++) {

int i = tid + idx * blockDim.x; if (i < cols) { float val = x_row[i]; if (val > max_val) max_val = val; }

}

同样的方式应用于指数和的计算以及最后的写入步骤。

这样，无论cols的大小如何，都可以处理。所以需要计算num_elements_per_thread。

综上，完整的CUDA代码可能如下：

首先包含必要的头文件：

#include <cuda_runtime.h> #include <math.h>

然后定义核函数：

global void softmax_kernel(float *d_x, float *d_y, int cols) {

extern __shared__ float shared_mem[]; float* max_shared = shared_mem; float* sum_shared = shared_mem + blockDim.x; int row = blockIdx.x; int tid = threadIdx.x; float *x_row = d_x + row * cols; float *y_row = d_y + row * cols; // 计算每行的最大值 float max_val = -INFINITY; int num_elements_per_thread = (cols + blockDim.x - 1) / blockDim.x; for (int idx = 0; idx < num_elements_per_thread; ++idx) { int i = tid + idx * blockDim.x; if (i < cols) { float val = x_row[i]; if (val > max_val) max_val = val; } } // 将线程的局部最大值存入共享内存，进行归约 max_shared[tid] = max_val; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { if (max_shared[tid + s] > max_shared[tid]) { max_shared[tid] = max_shared[tid + s]; } } __syncthreads(); } max_val = max_shared[0]; __syncthreads(); // 计算指数和 float exp_sum = 0.0f; for (int idx = 0; idx < num_elements_per_thread; ++idx) { int i = tid + idx * blockDim.x; if (i < cols) { exp_sum += expf(x_row[i] - max_val); } } // 归约求和 sum_shared[tid] = exp_sum; __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { sum_shared[tid] += sum_shared[tid + s]; } __syncthreads(); } exp_sum = sum_shared[0]; // 计算softmax并写入结果 if (exp_sum != 0.0f) { for (int idx = 0; idx < num_elements_per_thread; ++idx) { int i = tid + idx * blockDim.x; if (i < cols) { y_row[i] = expf(x_row[i] - max_val) / exp_sum; } } }

}

这里需要注意共享内存的使用，因为在同一个核函数中同时进行了最大值和求和的两个归约操作，所以共享内存需要有足够的空间。原来的代码中使用了一个共享内存数组，但最大值和求和的归约需要分开存储。因此，应该将共享内存分为两个部分，一个用于存储最大值归约的中间结果，另一个用于存储求和的中间结果。比如，共享内存的大小是2 * blockDim.x * sizeof(float)。所以在核函数启动时，需要分配足够的共享内存：

int threads_per_block = 256; size_t shared_mem_size = 2 * threads_per_block * sizeof(float); softmax_kernel<<<num_rows, threads_per_block, shared_mem_size>>>(d_x, d_y, cols);

这样，max_shared指向共享内存的前blockDim.x元素，sum_shared指向后blockDim.x元素。这样可以避免归约过程的数据覆盖问题。

但之前的代码中，我在共享内存的使用上可能有问题，因为当同时处理最大值和求和时，如果共享内存被重复使用而没有足够的空间，会导致数据覆盖。因此，这里必须确保最大值和求和的归约使用不同的共享内存区域。

这就是为什么需要将共享内存分割成两个部分：max_shared和 sum_shared，每个部分的大小是 blockDim.x 个元素。这样在归约时，不会互相干扰。

另外，在归约的过程中，比如在计算最大值时，每个线程把自己的最大值存入max_shared[tid]，然后通过树状归约得到全局最大值。同样的，在计算求和时，每个线程把局部的exp_sum存入sum_shared[tid]，然后归约。

这样修改后，共享内存的使用就不会有冲突了。

最后，在写入结果的时候，每个线程再次遍历所有自己负责的元素，计算exp(x_i - max_val)/exp_sum，并写入输出数组。

此外，需要考虑当exp_sum为0的情况，这时候可能所有元素都是负无穷大，导致指数和为0。这种情况可能需要特殊处理，但在实际应用中可能很少见，或者可以通过数值调整避免。这里假设不会发生，所以直接判断exp_sum是否为0，如果为0的话可能需要做一些处理，比如将结果设为0或者其他默认值。但示例中未处理这种情况，可能需要用户自己添加。

综上，这个核函数应该可以正确实现softmax的功能。不过还有优化的空间，例如使用更快的归约算法，或利用共享内存的bank冲突避免。此外，当处理非常大的行时，可能需要将数据分成多个块进行处理，或者使用动态并行性。但对于大多数情况，特别是行长度适中的情况，这个实现应该可行。

在调用该核函数的时候，需要确保每个线程块处理一行，并且线程块的数量等于输入数据的行数。比如，如果输入是形状为(rows, cols)的矩阵，则启动rows个线程块，每个线程块中有threads_per_block个线程，并分配足够共享内存。

此外，可能还需要在调用核函数之前检查输入的指针是否有效，以及CUDA错误处理等，但用户可能只需要核函数本身的代码。

最后，测试这个核函数时，应该用一些示例数据，比如向量[1.0, 2.0, 3.0]，计算softmax后的结果是否正确。例如，正确的结果应该是 [e^1/(e^1+e^2+e^3), e^2/(...), e^3/(...)]。数值计算过程是否正确，特别是减去最大值后的结果是否正确，比如上面的例子中最大值是3，所以每个元素减去3，得到[-2, -1, 0]，指数后的和为 e^-2 + e^-1 + e^0 ≈ 0.1353 + 0.3679 +1 = 1.5032，然后各自除以该和得到对应的概率。

<结束思考>

#include <cuda_runtime.h>#include <math.h>
__global__ void softmax_kernel(float* input, float* output, int num_rows, int num_cols) {extern __shared__ float shared_mem[];float* max_shared = shared_mem;float* sum_shared = shared_mem + blockDim.x;
int row = blockIdx.x;int tid = threadIdx.x;
if (row >= num_rows) return;
float* row_input = input + row * num_cols;float* row_output = output + row * num_cols;
// Step 1: 计算每行的最大值float max_val = -INFINITY;for (int i = tid; i < num_cols; i += blockDim.x) {max_val = fmaxf(max_val, row_input[i]);}
// 共享内存归约求最大值max_shared[tid] = max_val;__syncthreads();
for (int s = blockDim.x/2; s > 0; s >>= 1) {if (tid < s) {max_shared[tid] = fmaxf(max_shared[tid], max_shared[tid + s]);}__syncthreads();}max_val = max_shared[0];__syncthreads();
// Step 2: 计算指数和float exp_sum = 0.0f;for (int i = tid; i < num_cols; i += blockDim.x) {exp_sum += expf(row_input[i] - max_val);}
// 共享内存归约求指数和sum_shared[tid] = exp_sum;__syncthreads();
for (int s = blockDim.x/2; s > 0; s >>= 1) {if (tid < s) {sum_shared[tid] += sum_shared[tid + s];}__syncthreads();}exp_sum = sum_shared[0];__syncthreads();
// Step 3: 计算并写入Softmax结果if (exp_sum != 0.0f) {for (int i = tid; i < num_cols; i += blockDim.x) {row_output[i] = expf(row_input[i] - max_val) / exp_sum;}}}
void launch_softmax(float* input, float* output, int num_rows, int num_cols) {dim3 grid(num_rows);int block_size = 256; // 可根据实际情况调整线程块大小dim3 block(block_size);size_t shared_mem_size = 2 * block_size * sizeof(float); // 为归约分配足够的共享内存
softmax_kernel<<<grid, block, shared_mem_size>>>(input, output, num_rows, num_cols);}

import torcha = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])b = torch.softmax(a, 0)print(b)

# python testbench.pytensor([5.7661e-04, 1.5674e-03, 4.2606e-03, 1.1582e-02, 3.1482e-02, 8.5577e-02,2.3262e-01, 6.3233e-01])

#include <stdio.h>int main(void){float a[8] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};float b[8] = {0.0f};float * d_a, *d_b;cudaMalloc(&d_a, 8 * sizeof(float));cudaMalloc(&d_b, 8 * sizeof(float));cudaMemcpy(d_a, a, 8 *sizeof(float), cudaMemcpyHostToDevice);cudaMemset(d_b, 0, 8*sizeof(float));launch_softmax(d_a, d_b, 1, 8);cudaMemcpy(b, d_b, 8*sizeof(float), cudaMemcpyDeviceToHost);
for(int i = 0; i < 8; i ++){    printf("%e\n", b[i]);}cudaFree(d_a);cudaFree(d_b);return 0;}

计算结果：

# ./cuda_testbench5.766128e-041.567396e-034.260624e-031.158158e-023.148199e-028.557692e-022.326222e-016.323327e-01

CUDA 程序结果和 PyTorch 结果一致，说明正确性没问题！

接下来我们看 1.58 bit 量化版。它面对这个问题，足足思考了 4 分钟！

<思考过程>

好，用户让我写一段CUDA代码实现Softmax功能。我得先想想Softmax的公式，应该是对每个元素进行指数运算，然后除以所有元素指数的和。那在CUDA里，通常会用并行处理，尤其是对每个元素独立计算的部分。不过Softmax涉及到求和，可能需要同步或者原子操作，不过这样可能效率不高。

首先，我应该考虑如何分割输入数据。假设输入是一个一维的数组，那么每个线程可以处理一个元素。但是求和部分可能需要所有线程的合作，或者分成多个步骤。比如先计算每个元素的指数，然后计算总和，最后再每个元素除以总和。

不过，如果数组很大，直接求和可能会很困难，因为每个线程需要知道总和。这时候可能需要使用并行归约的方法来求总和。比如，先让每个线程计算自己的指数值，然后通过多级归约计算总和，然后再每个线程用这个总和来除以自己的指数值。

另外，内存管理也很重要。输入的数组和输出的数组需要在GPU上分配内存。还有，中间可能需要存储指数后的值，所以可能需要中间存储的空间。

那具体步骤可能是：

1. 每个线程计算自己的指数值，存到临时数组。

2. 计算所有指数值的总和。

3. 每个线程将自己的指数值除以总和，得到Softmax结果。

那代码的结构大概会是：先写一个内核函数，分成两个阶段。第一个内核计算指数并存到临时数组，然后调用另一个内核或者使用同一个内核的不同步骤来计算总和，再计算每个元素除以总和。

或者有没有可能在一个内核里完成？可能需要两次不同的内核调用。因为CUDA的同步机制有限，特别是跨块的同步，所以求和可能需要一个单独的内核或者使用CUDA的库函数来进行求和，比如使用Thrust或者cuBLAS，但用户可能希望纯CUDA代码。

或者考虑使用并行归约的方法，先让每个块计算自己部分的和，然后在第二个内核中将这些部分和相加，得到全局总和，然后每个元素再除以总和。不过这可能需要多次内核调用和额外的内存存储。

那具体步骤可能如下：

1. 分配输入和输出内存，以及中间存储的数组。

2. 第一个内核：每个线程计算对应元素的指数，并存到中间数组。

3. 调用一个归约内核，计算中间数组的总和。

4. 第二个内核：每个线程用中间数组的值除以总和，写入输出数组。

但这样的流程可能需要两次内核调用和一个中间数组，并且需要处理同步问题。或者有没有可能在一个内核中完成？不太可能，因为在计算总和的时候需要所有线程的贡献，这需要同步，而CUDA中块内的同步可以用__syncthreads()，但全局的同步做不到。

所以更实际的做法是将计算分成三个步骤：

1. 计算每个元素的指数并保存到临时数组。

2. 计算临时数组的总和。

3. 每个元素除以总和。

其中步骤2可能需要多次内核调用，比如使用并行归约的方法。

那现在具体代码如何编写呢？

比如，第一个内核计算指数：

global void compute_exp(float *input, float *temp, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { temp[idx] = exp(input[idx]); } }

然后第二个内核计算总和，可能需要一个归约函数：

global void sum_reduce(float *temp, float *sum, int size) { // 使用归约算法，每个块计算一部分的和，然后原子操作加到全局sum // 但具体实现可能需要多级处理，或者使用现有的归约方法 }

但归约可能比较复杂，可以考虑使用CUDA的库函数，比如用Thrust的sum，但用户希望纯CUDA代码的话，可能需要自己实现。

或者，可能更高效的方式是在同一个内核中，先让每个线程计算自己的指数，然后使用共享内存进行归约，计算总和。但是这只能在块级别进行，如果数组大小超过一个块，可能需要多级处理。

或者，或许用户输入的数组不是特别大，可以假设每个块处理一部分，然后分块计算每个块的总和，然后最后再汇总。这可能需要多个内核的步骤。

另一个方法是使用动态并行，但可能兼容性问题。

或者，可能更简单的方式是在计算完指数之后，将临时数组复制回主机，由主机计算总和，然后再传回GPU进行除法操作。但这样效率不高，尤其是对于GPU来说，但作为示例代码可能可以接受。

不过用户可能希望整个操作都在GPU上完成，所以需要纯GPU处理。

那可能的步骤：

• 分配输入、输出、临时数组，以及一个存储总和的变量。

• 第一个内核计算每个元素的指数，存入临时数组。: 计算exp(input[i])。. 然后，调用另一个内核来对这个临时数组求和，得到总和。

• 最后，调用第三个内核，每个线程用temp[i] / sum，得到输出。

现在，如何实现归约求和呢？可以参考标准的CUDA归约示例，比如从NVIDIA的官方归约例子。

例如，在归约内核中，使用共享内存来计算块内的部分和，然后写入到全局内存中的部分和数组，最后在另一个内核中将这些部分和相加得到总和。

不过这会增加代码的复杂度，但对于大规模数组来说，这是必要的。

如果数组较小，比如可以放在一个块里，那么可以用一个块来处理，在共享内存中计算总和。例如：

假设块的大小足够大，比如处理一个较小的数组，比如1024元素，一个块有1024线程，那么可以使用共享内存进行归约。

比如，在计算exp的同一个内核中，可能不行，因为每个线程计算自己的exp，然后归约总和。

或者，将计算exp和归约分开。

比如，首先调用compute_exp内核，然后调用一个归约内核来计算总和。

归约内核可能类似：

global void sum_kernel(float *temp, float *sum, int size) {shared float shared_data[1024]; int idx = blockIdx.x * blockDim.x + threadIdx.x; int tid = threadIdx.x;

// 每个线程加载数据到共享内存 if (idx < size) { shared_data[tid] = temp[idx]; } else { shared_data[tid] = 0; } __syncthreads(); // 进行归约，比如对1024线程，每次折半相加 for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { shared_data[tid] += shared_data[tid + s]; } __syncthreads(); } // 第一个线程将块的总和写入全局内存 if (tid == 0) { atomicAdd(sum, shared_data[0]); }

}

但这个方法只能在每个块处理一个部分，然后将每个块的总和通过原子操作加到全局sum中。如果数组很大，多个块会被启动，每个块计算自己的部分和，然后原子添加。

然后，当所有块执行完后，sum变量中存储的是整个数组的总和。然后，在第三步的内核中，每个线程用temp[idx] / sum。

但这里有几点需要注意：

• 原子操作可能会导致性能问题，但现代GPU的原子操作效率较高。"sum"变量应该是一个单个变量，用cudaMalloc分配，并在内核启动前初始化为0。

比如，在调用sum_kernel之前，先初始化sum为0：

float *d_sum; cudaMalloc(&d_sum, sizeof(float)); cudaMemset(d_sum, 0, sizeof(float));

然后在sum_kernel中，每个块将自己的部分和通过原子加到d_sum。

最后，将d_sum复制回主机，或者在第三个内核中直接使用。

但第三个内核需要访问总和值。或者，可以将总和存放在设备内存中，第三个内核可以直接读取。

不过，第三个内核启动时，sum已经被所有块处理完毕，因为CUDA内核调用是顺序执行的，所以在调用完sum_kernel之后，再调用第三个内核：

global void divide_kernel(float *temp, float *sum, float *output, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { output[idx] = temp[idx] / *sum; } }

这样，每个线程在divide_kernel中读取全局的sum值，除以每个temp元素。

不过，这可能存在一个问题：当sum_kernel中的原子操作完成后，所有块的处理完毕，sum的值是正确的总和。

不过，可能需要同步，或者确保在调用divide_kernel之前，sum_kernel已经完成。

这可以通过CUDA流来实现，但默认情况下，内核按顺序执行，所以只要sum_kernel和divide_kernel在同一个流，顺序执行，那么sum的值会被正确更新。

所以，整体步骤：

1. 分配输入、输出、temp、sum的GPU内存。

2. 将输入数据复制到GPU。

3. 调用compute_exp内核，计算temp数组。

4. 调用sum_kernel，计算sum的总和。

5. 调用divide_kernel，计算输出数组。

6. 将输出数组复制回主机。

但是sum_kernel需要每个块处理一部分，所以block和grid的划分需要根据temp数组的大小。

例如，假设temp数组大小为N，每个块处理最多1024元素，所以grid的大小是 (N + 1023) / 1024，每个block有1024 threads.

然后，在sum_kernel中，每个block处理自己的部分，使用原子操作将部分和加到sum变量中。

可能的问题：原子操作的正确性。当多个块同时执行原子添加，结果是否正确？是的，因为atomicAdd是按顺序进行的，所以最终sum会是所有部分的总和。

现在，把这些步骤转化为代码：

首先，初始化：

float *d_input, *d_temp, *d_output, *d_sum; cudaMalloc(&d_input, size * sizeof(float)); cudaMalloc(&d_temp, size * sizeof(float)); cudaMalloc(&d_output, size * sizeof(float)); cudaMalloc(&d_sum, sizeof(float));

然后，将输入数据从主机复制到d_input。

然后，调用compute_exp内核：

int blockSize = 256; int gridSize = (size + blockSize - 1) / blockSize; compute_exp<<<gridSize, blockSize>>>(d_input, d_temp, size);

接着，调用sum_kernel来求和：

cudaMemset(d_sum, 0, sizeof(float)); int sumBlockSize = 256; int sumGridSize = (size + sumBlockSize - 1) / sumBlockSize; sum_kernel<<<sumGridSize, sumBlockSize>>>(d_temp, d_sum, size);

这里sum_kernel的blockSize可能需要是2的幂，比如256，但共享内存的大小需要对应。比如，在sum_kernel中，共享内存的大小是256，所以需要声明为：

global void sum_kernel(float *temp, float *sum, int size) {shared float shared_data[256]; int idx = blockIdx.x * blockDim.x + threadIdx.x; int tid = threadIdx.x;

if (idx < size) { shared_data[tid] = temp[idx]; } else { shared_data[tid] = 0; } __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { shared_data[tid] += shared_data[tid + s]; } __syncthreads(); } if (tid == 0) { atomicAdd(sum, shared_data[0]); }

}

注意，每个块处理256个元素，然后将部分和加到sum中。

最后，调用divide_kernel：

divide_kernel<<<gridSize, blockSize>>>(d_temp, d_sum, d_output, size);

divide_kernel的实现：

global void divide_kernel(float *temp, float *sum, float *output, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { output[idx] = temp[idx] / *sum; } }

但是这里有个问题，在divide_kernel中，每个线程都要读取sum的值，而这个sum是一个指向单个float的指针，所以需要从设备内存中读取。这样可能会有效率问题，因为每个线程都会访问同一个全局内存地址，可能导致竞争，但实际情况下，只要sum已经被正确计算，并且所有线程读取的是同一个值，这是可行的。

不过，更高效的方法可能是将sum先复制到主机，然后再传回一个设备变量，或者使用常量。不过对于代码示例来说，可能没问题。

现在，需要考虑数值稳定性，比如在计算指数时可能出现溢出，特别是当输入值很大时。实际应用中，通常会减去最大值，来提高数值稳定性。但在用户的需求中，可能只是基本实现，所以可能不考虑这一点。但如果有需要，可以在计算指数前，找到数组的最大值，然后每个元素减去该最大值，再进行指数运算，这样可以避免指数太大导致溢出。例如，softmax的公式通常是：

softmax(x_i) = e^(x_i - max(x)) / sum(e^(x_j - max(x)))

但用户可能希望简单的实现，所以可能不包含这部分。

综上，完整的代码可能如下：

首先包含头文件：

#include <cuda_runtime.h> #include <stdio.h> #include <math.h>

然后定义三个内核函数：

global void compute_exp(float *input, float *temp, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { temp[idx] = expf(input[idx]); } }

global void sum_kernel(float *temp, float *sum, int size) {shared float shared_data[256]; int idx = blockIdx.x * blockDim.x + threadIdx.x; int tid = threadIdx.x;

if (idx < size) { shared_data[tid] = temp[idx]; } else { shared_data[tid] = 0.0f; } __syncthreads(); for (int s = blockDim.x / 2; s > 0; s >>= 1) { if (tid < s) { shared_data[tid] += shared_data[tid + s]; } __syncthreads(); } if (tid == 0) { atomicAdd(sum, shared_data[0]); }

}

global void divide_kernel(float *temp, float *sum, float *output, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { output[idx] = temp[idx] / *sum; } }

然后，主函数中的调用：

void cuda_softmax(float *h_input, float *h_output, int size) { float *d_input, *d_temp, *d_output, *d_sum;

cudaMalloc(&d_input, size * sizeof(float)); cudaMalloc(&d_temp, size * sizeof(float)); cudaMalloc(&d_output, size * sizeof(float)); cudaMalloc(&d_sum, sizeof(float)); cudaMemcpy(d_input, h_input, size * sizeof(float), cudaMemcpyHostToDevice); int blockSize = 256; int gridSize = (size + blockSize - 1) / blockSize; compute_exp<<<gridSize, blockSize>>>(d_input, d_temp, size); cudaMemset(d_sum, 0, sizeof(float)); int sumBlockSize = 256; int sumGridSize = (size + sumBlockSize - 1) / sumBlockSize; sum_kernel<<<sumGridSize, sumBlockSize>>>(d_temp, d_sum, size); divide_kernel<<<gridSize, blockSize>>>(d_temp, d_sum, d_output, size); cudaMemcpy(h_output, d_output, size * sizeof(float), cudaMemcpyDeviceToHost); cudaFree(d_input); cudaFree(d_temp); cudaFree(d_output); cudaFree(d_sum);

}

不过，这里可能存在几个问题：

1. sum_kernel中的原子操作是否正确。假设每个块计算其处理的部分的总和，然后原子加到d_sum。但是每个块处理的是多个元素，比如每个线程处理一个元素，但一个block有256线程，处理256元素。每个block计算该256元素的总和，然后原子添加到d_sum。当所有块处理完毕后，d_sum中的值就是所有元素的总和。

但是原sum_kernel的实现可能有问题。比如，当size不是256的倍数时，每个块处理的一部分可能超过实际元素的数量，导致shared_data中未被处理的线程位置被置为0，从而正确累加？

例如，idx < size时，线程读取temp[idx]，否则设置为0。这可能没问题，因为未处理的元素会被设为0，所以每个块的部分和仅针对实际存在的元素。

但每个块处理256个元素，如果总元素数不是256的倍数，最后一个块的部分会包含实际元素和多余的0。这样部分和的计算正确吗？例如，假设有257个元素，第一个块处理0-255，第二个块处理256-511，但只有256是实际存在，其余255个是超出范围，所以它们的shared_data[tid]会被设为0。所以在块的部分和中，第一个块的和是前256元素的exp的总和，第二个块的和是第257元素加上后面的255个0，总和正确。

但是，假设每个块处理256元素，当总元素数是N，每个块处理i256到i 256+255的元素。这样，每个块计算的部分和是这256个元素的总和（如果元素存在的话，否则为0），然后所有块的部分和相加，总和是正确的。

这样，sum_kernel应该正确计算总和。

另一个可能的优化是在sum_kernel中使用更高效的归约方法，比如每个块处理更多的元素，使用循环，但这样代码会更复杂。例如，每个线程处理多个元素，然后进行块内归约，这可能更适合处理更大的数组。

不过，作为示例代码，当前的方式可能足够。

另一个问题是，在divide_kernel中读取sum的值是否正确。因为sum是在设备内存中的一个变量，每个线程在divide_kernel中读取的是该变量的值，这个值已经被所有sum_kernel的块原子更新完成。在CUDA中，内核按照调用的顺序执行，所以在调用sum_kernel之后，调用divide_kernel，此时sum的值已经被完全更新。

不过，可能需要使用cudaDeviceSynchronize()来确保之前的内核执行完成，但在示例代码的主函数中，可能隐式调用了同步，因为CUDA的默认流是同步的，所以每个内核的调用是顺序执行的。

总的来说，这样的代码应该能够实现Softmax的功能，但可能不够高效或不够数值稳定。不过，作为示例代码，它应该能够展示基本的实现思路。

最后，测试一下代码，例如传入一个小的数组，比如h_input是[1.0, 2.0, 3.0]，计算每个元素的exp，总和是exp(1)+exp(2)+exp(3)，然后每个元素除以总和，输出是否正确。

可能存在一个问题：当多个线程同时读取sum时，是否会有同步问题？由于sum在设备内存中，并且所有写操作在sum_kernel中已经完成，divide_kernel中的读操作应该得到最终值。

综上，这应该是一个可行的CUDA实现，虽然可能有优化空间，比如处理数值稳定性问题，或者更高效的归约实现。

<结束思考>

#include <cuda_runtime.h>#include <math.h>
__global__ void compute_exp_and_sum(float *input, float *temp, float *sum, int size) {__shared__ float shared_data[256];int idx = blockIdx.x * blockDim.x + threadIdx.x;int tid = threadIdx.x;
// 计算指数并存储到临时数组if (idx < size) {temp[idx] = expf(input[idx]);shared_data[tid] = temp[idx];} else {shared_data[tid] = 0.0f;}
__syncthreads();
// 块内归约求和for (int s = blockDim.x / 2; s > 0; s >>= 1) {if (tid < s) {shared_data[tid] += shared_data[tid + s];}__syncthreads();}
// 原子操作将块的和加到全局总和if (tid == 0) {atomicAdd(sum, shared_data[0]);}}
__global__ void normalize(float *temp, float *sum, float *output, int size) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < size) {output[idx] = temp[idx] / *sum;}}
void cuda_softmax(float *h_input, float *h_output, int size) {float *d_input, *d_temp, *d_output, *d_sum;
// 分配设备内存cudaMalloc(&d_input, size * sizeof(float));cudaMalloc(&d_temp, size * sizeof(float));cudaMalloc(&d_output, size * sizeof(float));cudaMalloc(&d_sum, sizeof(float));
// 拷贝输入数据到设备cudaMemcpy(d_input, h_input, size * sizeof(float), cudaMemcpyHostToDevice);
int blockSize = 256;int gridSize = (size + blockSize - 1) / blockSize;
// 初始化总和为0cudaMemset(d_sum, 0, sizeof(float));
// 计算指数并求总和compute_exp_and_sum<<<gridSize, blockSize>>>(d_input, d_temp, d_sum, size);
// 归一化normalize<<<gridSize, blockSize>>>(d_temp, d_sum, d_output, size);
// 拷贝结果回主机cudaMemcpy(h_output, d_output, size * sizeof(float), cudaMemcpyDeviceToHost);
// 释放设备内存cudaFree(d_input);cudaFree(d_temp);cudaFree(d_output);cudaFree(d_sum);}

1.58bit 量化版相比满血版省掉了防指数累加溢出的代码，增加了显存分配和释放部分代码。编写和前面一样的测试用例如下：

#include <stdio.h>int main(void){float a[8] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};float b[8] = {0.0f};cuda_softmax(a, b, 8);for(int i = 0; i < 8; i ++){printf("%e\n", b[i]);}return 0;}

运行输出：

# ./cuda_testbench25.766128e-041.567396e-034.260624e-031.158158e-023.148199e-028.557692e-022.326222e-016.323327e-01

结果和 PyTorch 一致。

结论

• 条件允许时，能用满血版就用满血版，解决问题的专业性更强，能节省更多人力成本；

• 退而求其次，可选择量化版模型，思维链较完整，代码生成质量一般但可用；

• 如果实在没有条件，可以选择蒸馏版小模型，对硬件需求最低，但代码生成的可用性较差，需要付出更多时间和精力 debug；