forked from ccf-ai-infra/GPUCodeForces
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
# embeddingbag_cuda.py
|
|
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from embeddingbag_torch import BATCH_SIZE, EMB_DIM, VOCAB_SIZE # 导入维度常量
|
|
|
|
N_ELEMENTS = BATCH_SIZE * EMB_DIM
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
|
|
def __init__(self, embedding_weight):
|
|
super().__init__()
|
|
self.register_buffer('weight', embedding_weight)
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor embbag_forward_cuda(torch::Tensor input, torch::Tensor offsets, torch::Tensor weight);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
|
|
#define BLOCK_SIZE 256
|
|
|
|
__global__ void embbag_safe_kernel(
|
|
const long* input,
|
|
const long* offsets,
|
|
const float* weight,
|
|
float* output,
|
|
int num_batches, int feature_dim, int max_total_indices
|
|
) {
|
|
int batch_idx = blockIdx.x;
|
|
if (batch_idx >= num_batches) return;
|
|
|
|
int feature_idx = threadIdx.x;
|
|
if (feature_idx >= feature_dim) return;
|
|
|
|
long start_idx = offsets[batch_idx];
|
|
long end_idx = (batch_idx == num_batches - 1) ? max_total_indices : offsets[batch_idx + 1];
|
|
long seq_len = end_idx - start_idx;
|
|
|
|
float final_sum = 0.0f;
|
|
if (seq_len > 0) {
|
|
for (long i = start_idx; i < end_idx; ++i) {
|
|
long vocab_index = input[i];
|
|
final_sum += weight[vocab_index * feature_dim + feature_idx];
|
|
}
|
|
output[batch_idx * feature_dim + feature_idx] = final_sum / (float)seq_len;
|
|
} else {
|
|
output[batch_idx * feature_dim + feature_idx] = 0.0f; // Handle empty sequence
|
|
}
|
|
}
|
|
|
|
torch::Tensor embbag_forward_cuda(torch::Tensor input, torch::Tensor offsets, torch::Tensor weight) {
|
|
TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
|
|
TORCH_CHECK(input.dtype() == torch::kLong, "Input indices must be LongTensor");
|
|
TORCH_CHECK(offsets.dtype() == torch::kLong, "Offsets must be LongTensor");
|
|
|
|
input = input.contiguous();
|
|
offsets = offsets.contiguous();
|
|
weight = weight.contiguous();
|
|
|
|
int num_batches = offsets.size(0);
|
|
int feature_dim = weight.size(1);
|
|
int max_total_indices = input.size(0);
|
|
|
|
auto output = torch::empty({num_batches, feature_dim}, input.options().dtype(torch::kFloat32));
|
|
|
|
embbag_safe_kernel<<<num_batches, feature_dim>>>(
|
|
input.data_ptr<long>(),
|
|
offsets.data_ptr<long>(),
|
|
weight.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
num_batches, feature_dim, max_total_indices
|
|
);
|
|
|
|
return output;
|
|
}
|
|
|
|
"""
|
|
|
|
self.embbag_op = load_inline(
|
|
name="embbag_fused_op_safe_final_v2",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["embbag_forward_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=True
|
|
)
|
|
|
|
def forward(self, input: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
|
|
return self.embbag_op.embbag_forward_cuda(input, offsets, self.weight) |