forked from ccf-ai-infra/GPUCodeForces
235 lines
8.7 KiB
Python
235 lines
8.7 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.cpp_extension import load_inline
|
||
# 从 torch 文件导入常量
|
||
from infonceloss_torch import BATCH_SIZE, FEATURE_DIM, TEMPERATURE, N_NEGATIVES
|
||
|
||
|
||
class ModelNew(nn.Module):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.temperature = TEMPERATURE
|
||
self._compile_cuda_kernel()
|
||
|
||
def _compile_cuda_kernel(self):
|
||
cpp_source = """
|
||
#include <torch/extension.h>
|
||
|
||
// C++ 接口
|
||
torch::Tensor infonce_forward_cuda(
|
||
torch::Tensor query, // (B, D)
|
||
torch::Tensor positive, // (B, D)
|
||
torch::Tensor negative_sims, // (B, N) - 预先计算的
|
||
float temperature
|
||
);
|
||
"""
|
||
|
||
cuda_source = """
|
||
#include <torch/extension.h>
|
||
#include <cuda_runtime.h>
|
||
#include <cmath>
|
||
#include <float.h> // For FLT_MAX
|
||
|
||
// 使用 256 个线程的块大小
|
||
#define BLOCK_SIZE 256
|
||
|
||
/*
|
||
* InfoNCE 融合核函数
|
||
* 我们启动 B 个块 (gridDim.x = B)，每个块负责一行 (一个 query) 的 loss 计算。
|
||
* 每个块 (blockIdx.x) 计算:
|
||
* 1. query[i] 和 positive[i] 之间的点积 (pos_logit)
|
||
* 2. 对 [pos_logit, neg_logits[i,:]] 执行稳定的 LogSumExp
|
||
* 3. 计算 loss_i = -pos_logit + logsumexp
|
||
*
|
||
* @param loss_per_row_out - (B,) 形状的张量，用于存储 loss_i
|
||
*/
|
||
__global__ void infonce_fused_kernel(
|
||
const float* __restrict__ query_data, // (B, D)
|
||
const float* __restrict__ positive_data, // (B, D)
|
||
const float* __restrict__ negative_sims_data, // (B, N)
|
||
float* __restrict__ loss_per_row_out, // (B,)
|
||
int B,
|
||
int D,
|
||
int N,
|
||
float temperature
|
||
) {
|
||
// 每个块计算一行
|
||
int i = blockIdx.x; // 当前 query 的索引 (0 到 B-1)
|
||
if (i >= B) return;
|
||
|
||
// --- 共享内存 ---
|
||
// s_dot 用于计算 pos_logit
|
||
__shared__ float s_dot[BLOCK_SIZE];
|
||
// s_max 和 s_sum 用于稳定的 LogSumExp
|
||
__shared__ float s_max[BLOCK_SIZE];
|
||
__shared__ float s_sum[BLOCK_SIZE];
|
||
|
||
// --- 1. 计算 Positive Logit ---
|
||
// 融合了 F.cosine_similarity(query[i], positive[i]) / temp
|
||
float thread_dot_sum = 0.0f;
|
||
|
||
// 使用 Grid-Stride 循环计算点积
|
||
for (int k = threadIdx.x; k < D; k += BLOCK_SIZE) {
|
||
thread_dot_sum += query_data[i * D + k] * positive_data[i * D + k];
|
||
}
|
||
s_dot[threadIdx.x] = thread_dot_sum;
|
||
|
||
// 块内归约 (Sum)
|
||
__syncthreads();
|
||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
||
if (threadIdx.x < offset) {
|
||
s_dot[threadIdx.x] += s_dot[threadIdx.x + offset];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
// 线程 0 现在拥有 pos_logit
|
||
// 我们将其存储在 s_dot[0] 中以供后续步骤使用
|
||
if (threadIdx.x == 0) {
|
||
s_dot[0] = s_dot[0] / temperature;
|
||
}
|
||
__syncthreads(); // 确保所有线程都能读到 s_dot[0]
|
||
|
||
const float pos_logit = s_dot[0]; // 所有线程的常量
|
||
|
||
// --- 2. 稳定的 LogSumExp (Pass 1: Find Max) ---
|
||
float thread_max = -FLT_MAX;
|
||
|
||
// 线程 0 包含 pos_logit
|
||
if (threadIdx.x == 0) {
|
||
thread_max = pos_logit;
|
||
}
|
||
|
||
// 遍历 N 个 negative logits
|
||
for (int j = threadIdx.x; j < N; j += BLOCK_SIZE) {
|
||
float neg_logit = negative_sims_data[i * N + j] / temperature;
|
||
thread_max = fmaxf(thread_max, neg_logit);
|
||
}
|
||
s_max[threadIdx.x] = thread_max;
|
||
|
||
// 块内归约 (Max)
|
||
__syncthreads();
|
||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
||
if (threadIdx.x < offset) {
|
||
s_max[threadIdx.x] = fmaxf(s_max[threadIdx.x], s_max[threadIdx.x + offset]);
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
// 线程 0 拥有 global_max
|
||
if (threadIdx.x == 0) {
|
||
s_max[0] = s_max[0];
|
||
}
|
||
__syncthreads(); // 确保所有线程都能读到 s_max[0]
|
||
|
||
const float global_max = s_max[0];
|
||
|
||
// --- 3. 稳定的 LogSumExp (Pass 2: Sum Exp Diff) ---
|
||
float thread_sum_exp = 0.0f;
|
||
|
||
// 线程 0 添加 positive_logit 的贡献
|
||
if (threadIdx.x == 0) {
|
||
thread_sum_exp = expf(pos_logit - global_max);
|
||
}
|
||
|
||
// 遍历 N 个 negative logits
|
||
for (int j = threadIdx.x; j < N; j += BLOCK_SIZE) {
|
||
float neg_logit = negative_sims_data[i * N + j] / temperature;
|
||
thread_sum_exp += expf(neg_logit - global_max);
|
||
}
|
||
s_sum[threadIdx.x] = thread_sum_exp;
|
||
|
||
// 块内归约 (Sum)
|
||
__syncthreads();
|
||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
||
if (threadIdx.x < offset) {
|
||
s_sum[threadIdx.x] += s_sum[threadIdx.x + offset];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
// --- 4. 计算最终的 loss[i] ---
|
||
if (threadIdx.x == 0) {
|
||
float log_sum_exp = global_max + logf(s_sum[0]);
|
||
// loss_i = -logit[0] + logsumexp
|
||
float loss_i = -pos_logit + log_sum_exp;
|
||
loss_per_row_out[i] = loss_i;
|
||
}
|
||
}
|
||
|
||
// C++ 封装函数
|
||
torch::Tensor infonce_forward_cuda(
|
||
torch::Tensor query,
|
||
torch::Tensor positive,
|
||
torch::Tensor negative_sims, // 注意：这是未缩放的
|
||
float temperature
|
||
) {
|
||
// 检查
|
||
TORCH_CHECK(query.is_cuda(), "query must be a CUDA tensor");
|
||
TORCH_CHECK(positive.is_cuda(), "positive must be a CUDA tensor");
|
||
TORCH_CHECK(negative_sims.is_cuda(), "negative_sims must be a CUDA tensor");
|
||
|
||
query = query.contiguous();
|
||
positive = positive.contiguous();
|
||
negative_sims = negative_sims.contiguous();
|
||
|
||
const int B = query.size(0);
|
||
const int D = query.size(1);
|
||
const int N = negative_sims.size(1);
|
||
|
||
TORCH_CHECK(positive.size(0) == B && positive.size(1) == D, "positive tensor has wrong size");
|
||
TORCH_CHECK(negative_sims.size(0) == B, "negative_sims tensor has wrong size");
|
||
|
||
// 分配一个张量来保存每个块 (每行) 的 loss
|
||
auto loss_per_row = torch::empty({B}, query.options());
|
||
|
||
const int block_size = BLOCK_SIZE;
|
||
const int grid_size = B; // B 个块，每个块处理一行
|
||
|
||
// 启动 CUDA 核函数
|
||
infonce_fused_kernel<<<grid_size, block_size>>>(
|
||
query.data_ptr<float>(),
|
||
positive.data_ptr<float>(),
|
||
negative_sims.data_ptr<float>(),
|
||
loss_per_row.data_ptr<float>(),
|
||
B, D, N,
|
||
temperature
|
||
);
|
||
|
||
|
||
// 核函数返回后，loss_per_row 包含 B 个 loss 值
|
||
// 我们需要对它们取平均
|
||
return loss_per_row.mean();
|
||
}
|
||
"""
|
||
|
||
# JIT (Just-In-Time) 编译
|
||
self.infonce_op = load_inline(
|
||
name="infonce_op_v1_stable",
|
||
cpp_sources=cpp_source,
|
||
cuda_sources=cuda_source,
|
||
functions=["infonce_forward_cuda"],
|
||
extra_cuda_cflags=["-O3"],
|
||
verbose=False
|
||
)
|
||
|
||
def forward(self, query: torch.Tensor, positive: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
|
||
# 1. (Python) 执行优化的 matmul (cuBLAS)
|
||
# (B, D) @ (D, N) -> (B, N)
|
||
# 这是未缩放的 (没有 / temp)
|
||
negative_sims_unscaled = torch.matmul(query, negatives.t())
|
||
|
||
# 2. (CUDA) 调用融合核函数
|
||
# 核函数将处理:
|
||
# - query, positive 的 cosine similarity
|
||
# - 对所有 sim 应用 / temp
|
||
# - 稳定的 LogSumExp 和 CrossEntropy
|
||
# - 最终的 Mean 归约
|
||
return self.infonce_op.infonce_forward_cuda(
|
||
query,
|
||
positive,
|
||
negative_sims_unscaled,
|
||
self.temperature
|
||
) |