Softmax 是深度学习模型中常见的归一化算子，广泛用于分类输出、Transformer Attention 权重计算等场景。本文围绕 Softmax 算子展开，重点学习 HIP 中的 block 内归约、shared memory 使用、数值稳定性处理，以及当 `cols` 较大时如何通过多 block 协作提升并行度。 ## 1. Softmax算子定义 Softmax 将一组实数映射为概率分布，其定义为： $$ y_i = \frac{e^{x_i}}{\sum_j e^{x_j}} $$ 在实际工程实现中，为了避免 `exp(x)` 在输入较大时发生数值溢出，通常采用数值稳定版本： $$ y_i = \frac{e^{x_i - max(x)}}{\sum_j e^{x_j - max(x)}} $$ 也就是先求当前行的最大值，再计算 `exp(x - max)`，最后除以指数和。 假设输入为二维矩阵： ```text input: [rows, cols] output: [rows, cols] ``` 本文中 Softmax 沿最后一个维度 `cols` 计算，即每一行单独做归一化。 示例： ```text input: [ [1, 2, 3], [1, 1, 1] ] output: [ [0.0900, 0.2447, 0.6652], [0.3333, 0.3333, 0.3333] ] ``` ## 2. Softmax并行计算思路 Softmax 对每一行数据的计算主要包括三个步骤： 1. 求当前行最大值 `max_val` 2. 计算 `exp(x - max_val)` 并求指数和 `sum_val` 3. 将每个元素除以 `sum_val` 得到归一化结果 对于 `cols` 不特别大的情况，可以采用： - 一个 block 处理一行数据； - block 内多个 thread 并行处理该行的不同列； - 使用 shared memory 完成 block 内最大值归约和求和归约； - 每个线程按步长 `blockDim.x` 处理多个元素。 这种方式结构清晰，适合教学和中小规模 `cols` 场景，例如 `cols <= 1024` 或几千以内的行向量。 ## 3. 基础版Softmax Kernel 基础版 Kernel 使用一个 block 处理一行数据，并通过动态 shared memory 存储每个线程的局部最大值和局部求和值。 ```cpp #include #include #include __global__ void softmax_kernel(const float* input, float* output, int rows, int cols) { extern __shared__ float shared[]; int row = blockIdx.x; int tid = threadIdx.x; if (row >= rows) { return; } const float* row_input = input + row * cols; float* row_output = output + row * cols; float local_max = -FLT_MAX; for (int col = tid; col < cols; col += blockDim.x) { local_max = fmaxf(local_max, row_input[col]); } shared[tid] = local_max; __syncthreads(); for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] = fmaxf(shared[tid], shared[tid + stride]); } __syncthreads(); } float max_val = shared[0]; float local_sum = 0.0f; for (int col = tid; col < cols; col += blockDim.x) { float e = expf(row_input[col] - max_val); row_output[col] = e; local_sum += e; } shared[tid] = local_sum; __syncthreads(); for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] += shared[tid + stride]; } __syncthreads(); } float sum_val = shared[0]; for (int col = tid; col < cols; col += blockDim.x) { row_output[col] /= sum_val; } } ``` ## 4. Host端调用封装 Host 端负责分配显存、拷贝输入、配置 kernel 启动参数、同步并拷贝结果。 ```cpp #include #include #include #define HIP_CHECK(cmd) \ do { \ hipError_t error = cmd; \ if (error != hipSuccess) { \ std::cerr << "HIP error: " << hipGetErrorString(error) \ << " at " << __FILE__ << ":" << __LINE__ \ << std::endl; \ std::exit(EXIT_FAILURE); \ } \ } while (0) void softmax_hip(const float* h_input, float* h_output, int rows, int cols) { float* d_input = nullptr; float* d_output = nullptr; size_t size = static_cast(rows) * cols * sizeof(float); HIP_CHECK(hipMalloc(&d_input, size)); HIP_CHECK(hipMalloc(&d_output, size)); HIP_CHECK(hipMemcpy(d_input, h_input, size, hipMemcpyHostToDevice)); int threads = 256; if (cols < threads) { threads = 1; while (threads < cols) { threads <<= 1; } } dim3 block(threads); dim3 grid(rows); size_t shared_mem_size = threads * sizeof(float); hipLaunchKernelGGL( softmax_kernel, grid, block, shared_mem_size, 0, d_input, d_output, rows, cols ); HIP_CHECK(hipGetLastError()); HIP_CHECK(hipDeviceSynchronize()); HIP_CHECK(hipMemcpy(h_output, d_output, size, hipMemcpyDeviceToHost)); HIP_CHECK(hipFree(d_input)); HIP_CHECK(hipFree(d_output)); } ``` 其中线程数选择策略如下： | `cols`规模 | 推荐线程数策略 | | ------------- | ------------------------------------- | | `cols < 256` | 使用不小于 `cols` 的 2 的幂 | | `cols >= 256` | 使用 256 个线程，每个线程处理多个元素 | | `cols` 很大 | 使用多 block 处理一行，见第 7 节 | ## 5. Kernel关键优化点 ### 一个block处理一行 ```cpp int row = blockIdx.x; ``` 每个 block 负责输入矩阵的一行。这种设计简单直观，适合二维张量最后一维 Softmax。 ### 使用shared memory做归约 ```cpp extern __shared__ float shared[]; ``` 启动 kernel 时动态分配 shared memory： ```cpp size_t shared_mem_size = threads * sizeof(float); ``` shared memory 用于存放每个线程的局部最大值和局部求和值，并在 block 内完成归约。 ### 保证数值稳定性 直接计算 `expf(x)` 时，如果 `x` 较大，容易发生溢出。稳定 Softmax 使用： ```cpp expf(x - max_val) ``` 由于 `x - max_val <= 0`，可以显著降低指数计算的溢出风险。 ### 使用步长循环处理长行 当 `cols > blockDim.x` 时，每个线程需要处理多个元素： ```cpp for (int col = tid; col < cols; col += blockDim.x) ``` 这种写法能够让固定数量的线程覆盖整行数据。 ## 6. 大cols场景下的Softmax实现 当 `cols` 很大时，例如： ```text cols = 8192 / 16384 / 32768 / 65536 ``` 一个 block 处理一整行会遇到以下问题： 1. 单个 block 线程数有限，通常不能超过 1024； 2. 每个线程需要处理很多元素，单行并行度不足； 3. 如果 `rows` 较小，GPU 上可调度的 block 数很少，设备利用率偏低； 4. `__syncthreads()` 只能同步同一个 block 内线程，不能直接跨 block 同步。 因此，大 `cols` 场景通常使用多个 block 处理同一行，并通过多个 kernel launch 形成全局同步边界。 ### 整体流程 大 `cols` 版本可以拆成 5 个阶段： ```text 每行分 tile 求 partial max 对每行 partial max 再归约，得到 row_max 计算 exp(input - row_max)，并求 partial sum 对每行 partial sum 再归约，得到 row_sum 根据 row_sum 完成归一化 ``` 设定： ```cpp #define BLOCK_SIZE 256 #define ITEMS_PER_THREAD 4 ``` 则每个 block 处理： ```text 256 * 4 = 1024 ``` 个元素。若 `cols = 65536`，则每行会被拆成 64 个 tile。 ### 局部最大值Kernel ```cpp #define BLOCK_SIZE 256 #define ITEMS_PER_THREAD 4 __global__ void softmax_partial_max_kernel(const float* input, float* partial_max, int rows, int cols, int num_tiles) { extern __shared__ float shared[]; int tile = blockIdx.x; int row = blockIdx.y; int tid = threadIdx.x; if (row >= rows || tile >= num_tiles) { return; } int tile_size = BLOCK_SIZE * ITEMS_PER_THREAD; int tile_start = tile * tile_size; const float* row_input = input + row * cols; float local_max = -FLT_MAX; #pragma unroll for (int i = 0; i < ITEMS_PER_THREAD; ++i) { int col = tile_start + tid + i * BLOCK_SIZE; if (col < cols) { local_max = fmaxf(local_max, row_input[col]); } } shared[tid] = local_max; __syncthreads(); for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] = fmaxf(shared[tid], shared[tid + stride]); } __syncthreads(); } if (tid == 0) { partial_max[row * num_tiles + tile] = shared[0]; } } ``` ### 每行最大值归约Kernel ```cpp __global__ void softmax_final_max_kernel(const float* partial_max, float* row_max, int rows, int num_tiles) { extern __shared__ float shared[]; int row = blockIdx.x; int tid = threadIdx.x; if (row >= rows) { return; } float local_max = -FLT_MAX; for (int tile = tid; tile < num_tiles; tile += blockDim.x) { local_max = fmaxf(local_max, partial_max[row * num_tiles + tile]); } shared[tid] = local_max; __syncthreads(); for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] = fmaxf(shared[tid], shared[tid + stride]); } __syncthreads(); } if (tid == 0) { row_max[row] = shared[0]; } } ``` ### 指数计算与局部求和Kernel ```cpp __global__ void softmax_partial_sum_kernel(const float* input, float* output, const float* row_max, float* partial_sum, int rows, int cols, int num_tiles) { extern __shared__ float shared[]; int tile = blockIdx.x; int row = blockIdx.y; int tid = threadIdx.x; if (row >= rows || tile >= num_tiles) { return; } int tile_size = BLOCK_SIZE * ITEMS_PER_THREAD; int tile_start = tile * tile_size; const float* row_input = input + row * cols; float* row_output = output + row * cols; float max_val = row_max[row]; float local_sum = 0.0f; #pragma unroll for (int i = 0; i < ITEMS_PER_THREAD; ++i) { int col = tile_start + tid + i * BLOCK_SIZE; if (col < cols) { float e = expf(row_input[col] - max_val); row_output[col] = e; local_sum += e; } } shared[tid] = local_sum; __syncthreads(); for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] += shared[tid + stride]; } __syncthreads(); } if (tid == 0) { partial_sum[row * num_tiles + tile] = shared[0]; } } ``` ### 每行求和归约Kernel ```cpp __global__ void softmax_final_sum_kernel(const float* partial_sum, float* row_sum, int rows, int num_tiles) { extern __shared__ float shared[]; int row = blockIdx.x; int tid = threadIdx.x; if (row >= rows) { return; } float local_sum = 0.0f; for (int tile = tid; tile < num_tiles; tile += blockDim.x) { local_sum += partial_sum[row * num_tiles + tile]; } shared[tid] = local_sum; __syncthreads(); for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { if (tid < stride) { shared[tid] += shared[tid + stride]; } __syncthreads(); } if (tid == 0) { row_sum[row] = shared[0]; } } ``` ### 归一化Kernel ```cpp __global__ void softmax_normalize_kernel(float* output, const float* row_sum, int rows, int cols) { int row = blockIdx.y; int tile = blockIdx.x; int tid = threadIdx.x; if (row >= rows) { return; } int tile_size = BLOCK_SIZE * ITEMS_PER_THREAD; int tile_start = tile * tile_size; float* row_output = output + row * cols; float sum_val = row_sum[row]; #pragma unroll for (int i = 0; i < ITEMS_PER_THREAD; ++i) { int col = tile_start + tid + i * BLOCK_SIZE; if (col < cols) { row_output[col] /= sum_val; } } } ``` ## 7. 大cols版本Host端封装 大 `cols` 版本需要额外分配 `partial_max`、`partial_sum`、`row_max` 和 `row_sum` 等中间缓冲区。 ```cpp void softmax_hip_large_cols(const float* h_input, float* h_output, int rows, int cols) { float* d_input = nullptr; float* d_output = nullptr; float* d_partial_max = nullptr; float* d_partial_sum = nullptr; float* d_row_max = nullptr; float* d_row_sum = nullptr; size_t data_size = static_cast(rows) * cols * sizeof(float); int tile_size = BLOCK_SIZE * ITEMS_PER_THREAD; int num_tiles = (cols + tile_size - 1) / tile_size; size_t partial_size = static_cast(rows) * num_tiles * sizeof(float); size_t row_size = static_cast(rows) * sizeof(float); HIP_CHECK(hipMalloc(&d_input, data_size)); HIP_CHECK(hipMalloc(&d_output, data_size)); HIP_CHECK(hipMalloc(&d_partial_max, partial_size)); HIP_CHECK(hipMalloc(&d_partial_sum, partial_size)); HIP_CHECK(hipMalloc(&d_row_max, row_size)); HIP_CHECK(hipMalloc(&d_row_sum, row_size)); HIP_CHECK(hipMemcpy(d_input, h_input, data_size, hipMemcpyHostToDevice)); dim3 block(BLOCK_SIZE); dim3 grid_tiles(num_tiles, rows); dim3 grid_rows(rows); size_t shared_size = BLOCK_SIZE * sizeof(float); hipLaunchKernelGGL( softmax_partial_max_kernel, grid_tiles, block, shared_size, 0, d_input, d_partial_max, rows, cols, num_tiles ); HIP_CHECK(hipGetLastError()); hipLaunchKernelGGL( softmax_final_max_kernel, grid_rows, block, shared_size, 0, d_partial_max, d_row_max, rows, num_tiles ); HIP_CHECK(hipGetLastError()); hipLaunchKernelGGL( softmax_partial_sum_kernel, grid_tiles, block, shared_size, 0, d_input, d_output, d_row_max, d_partial_sum, rows, cols, num_tiles ); HIP_CHECK(hipGetLastError()); hipLaunchKernelGGL( softmax_final_sum_kernel, grid_rows, block, shared_size, 0, d_partial_sum, d_row_sum, rows, num_tiles ); HIP_CHECK(hipGetLastError()); hipLaunchKernelGGL( softmax_normalize_kernel, grid_tiles, block, 0, 0, d_output, d_row_sum, rows, cols ); HIP_CHECK(hipGetLastError()); HIP_CHECK(hipDeviceSynchronize()); HIP_CHECK(hipMemcpy(h_output, d_output, data_size, hipMemcpyDeviceToHost)); HIP_CHECK(hipFree(d_input)); HIP_CHECK(hipFree(d_output)); HIP_CHECK(hipFree(d_partial_max)); HIP_CHECK(hipFree(d_partial_sum)); HIP_CHECK(hipFree(d_row_max)); HIP_CHECK(hipFree(d_row_sum)); } ``` 可以根据 `cols` 大小选择不同实现： ```cpp void softmax_hip_auto(const float* h_input, float* h_output, int rows, int cols) { if (cols <= 4096) { softmax_hip(h_input, h_output, rows, cols); } else { softmax_hip_large_cols(h_input, h_output, rows, cols); } } ``` 阈值 `4096` 不是固定值，应根据 DCU 型号、数据规模、batch size 和实际 benchmark 结果调整。 ## 8. 编译与运行 在 DCU 环境中，通常需要先确认 HIP/DTK 环境是否可用： ```bash hipcc --version rocminfo ``` ### 基础版程序编译运行： ```bash hipcc softmax_basic.cpp -o softmax_basic ./softmax_basic ``` 可能输出如下： ```text HIP Softmax output: 0.0320586 0.0871443 0.236883 0.643914 0.25 0.25 0.25 0.25 CPU reference: 0.0320586 0.0871443 0.236883 0.643914 0.25 0.25 0.25 0.25 ``` ### 大 `cols` 版本编译运行： ```bash hipcc softmax_large_cols.cpp -o softmax_large_cols ./softmax_large_cols ``` 可通过最大误差检查结果： ```text Large-cols HIP Softmax test rows = 4, cols = 32768, tile_size = 1024, num_tiles = 32 max error = 3.1141099e-09 max row sum error = 3.528595e-05 first row first 8 outputs: 4.7680678e-06 4.9194227e-06 5.0755821e-06 5.2366981e-06 5.4029288e-06 5.5744363e-06 5.751388e-06 5.9339563e-06 Result check passed. ``` 不同硬件、编译器版本和输入规模下，误差和运行时间可能略有差异。通常可以使用 `1e-4` 量级作为单精度浮点结果校验阈值。