diff --git a/README.md b/README.md index 4785a3b..417fbd1 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,9 @@ 该清单由 `scripts/update_operator_checklist.py` 根据 `scripts/operator_targets.txt` 和 `S1 codes/` 自动生成；新增待实现算子请写入目标清单，新增实现目录后运行 `python scripts/update_operator_checklist.py --sync-targets` 即可自动勾选。 -- 已实现：577 +- 已实现：578 - 未实现：0 -- 跟踪总数：577 +- 跟踪总数：578
展开查看算子实现状态 @@ -231,6 +231,7 @@ | [x] | gaussian_pdf | [ZZZJ_#120](S1%20codes/ZZZJ_%23120) | | [x] | GaussianNLLLoss | [hli28146_#113](S1%20codes/hli28146_%23113)
[uucoco 31](S1%20codes/uucoco%2031)
[uucoco_#1](S1%20codes/uucoco_%231) | | [x] | GDL | [hli28146_#55](S1%20codes/hli28146_%2355) | +| [x] | geglu | [geglu_sample](S1%20codes/geglu_sample) | | [x] | GELU-Affine-Gate | [Ljy123_#94](S1%20codes/Ljy123_%2394) | | [x] | gelu_dropout | [Ljy123_#3](S1%20codes/Ljy123_%233) | | [x] | gempool | [hli28146_#12](S1%20codes/hli28146_%2312) | diff --git a/S1 codes/geglu_sample/geglu_cuda.py b/S1 codes/geglu_sample/geglu_cuda.py new file mode 100644 index 0000000..51412ab --- /dev/null +++ b/S1 codes/geglu_sample/geglu_cuda.py @@ -0,0 +1,89 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + + +class ModelNew(nn.Module): + def __init__(self): + super().__init__() + self.op = load_inline( + name="geglu_fused_tanh_v1", + cpp_sources=""" + #include + torch::Tensor geglu_cuda(torch::Tensor input); + """, + cuda_sources=""" + #include + #include + + __device__ __forceinline__ float gelu_tanh(float x) { + const float kAlpha = 0.7978845608028654f; + const float kBeta = 0.044715f; + float x3 = x * x * x; + float inner = kAlpha * (x + kBeta * x3); + return 0.5f * x * (1.0f + tanhf(inner)); + } + + __global__ void geglu_kernel( + const float* __restrict__ input, + float* __restrict__ output, + long long output_numel, + int hidden, + int input_last_dim + ) { + long long idx = blockIdx.x * blockDim.x + threadIdx.x; + long long stride = (long long)blockDim.x * gridDim.x; + + for (long long i = idx; i < output_numel; i += stride) { + int col = i % hidden; + long long row = i / hidden; + long long base = row * input_last_dim + col; + float value = input[base]; + float gate = input[base + hidden]; + output[i] = value * gelu_tanh(gate); + } + } + + torch::Tensor geglu_cuda(torch::Tensor input) { + TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); + TORCH_CHECK(input.scalar_type() == torch::kFloat32, "input must be float32"); + TORCH_CHECK(input.dim() >= 1, "input must have at least one dimension"); + + auto x = input.contiguous(); + int input_last_dim = x.size(-1); + TORCH_CHECK(input_last_dim % 2 == 0, "last dimension must be even for GEGLU"); + + int hidden = input_last_dim / 2; + auto out_sizes = x.sizes().vec(); + out_sizes.back() = hidden; + auto output = torch::empty(out_sizes, x.options()); + + long long output_numel = output.numel(); + if (output_numel == 0) { + return output; + } + + int threads = 256; + int blocks = (int)((output_numel + threads - 1) / threads); + if (blocks > 65535) { + blocks = 65535; + } + + geglu_kernel<<>>( + x.data_ptr(), + output.data_ptr(), + output_numel, + hidden, + input_last_dim + ); + + return output; + } + """, + functions=["geglu_cuda"], + extra_cuda_cflags=["-O3", "--use_fast_math"], + verbose=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.op.geglu_cuda(x) diff --git a/S1 codes/geglu_sample/geglu_torch.py b/S1 codes/geglu_sample/geglu_torch.py new file mode 100644 index 0000000..68641fe --- /dev/null +++ b/S1 codes/geglu_sample/geglu_torch.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class Model(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + value, gate = x.chunk(2, dim=-1) + return value * F.gelu(gate, approximate="tanh") + + +batch_size = 2048 +feature_dim = 4096 + + +def get_inputs(): + x = torch.randn(batch_size, feature_dim, dtype=torch.float32) * 3.0 + return [x] + + +def get_init_inputs(): + return [] diff --git a/S1 codes/geglu_sample/prompt.txt b/S1 codes/geglu_sample/prompt.txt new file mode 100644 index 0000000..547fb35 --- /dev/null +++ b/S1 codes/geglu_sample/prompt.txt @@ -0,0 +1,15 @@ +Operator: GEGLU + +Implement a fused CUDA kernel for the GEGLU activation used in Transformer MLP blocks. + +Reference PyTorch behavior: + +```python +value, gate = x.chunk(2, dim=-1) +y = value * torch.nn.functional.gelu(gate, approximate="tanh") +``` + +The input is a contiguous or non-contiguous float32 CUDA tensor whose last +dimension is even. The output keeps the same leading dimensions and halves the +last dimension. The CUDA implementation should fuse the chunk, tanh-approx GELU, +and elementwise multiply into a single pass over output elements. diff --git a/S1 codes/geglu_sample/run_code.py b/S1 codes/geglu_sample/run_code.py new file mode 100644 index 0000000..95e4a4b --- /dev/null +++ b/S1 codes/geglu_sample/run_code.py @@ -0,0 +1,67 @@ +import time + +import torch + +from geglu_cuda import ModelNew +from geglu_torch import Model, get_init_inputs, get_inputs + + +def _to_cuda(values): + return [x.cuda() if isinstance(x, torch.Tensor) else x for x in values] + + +def run_benchmark(): + if not torch.cuda.is_available(): + print("CUDA is not available.") + return False, 0.0 + + init_inputs = _to_cuda(get_init_inputs()) + inputs = _to_cuda(get_inputs()) + + torch_model = Model(*init_inputs).cuda().eval() + cuda_model = ModelNew(*init_inputs).cuda().eval() + + with torch.no_grad(): + output_torch = torch_model(*inputs) + output_cuda = cuda_model(*inputs) + + max_diff = (output_torch - output_cuda).abs().max().item() + mean_diff = (output_torch - output_cuda).abs().mean().item() + precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-4, atol=1e-4) + + print("-------------------- precision check --------------------") + print(f"max diff: {max_diff:.8f}") + print(f"mean diff: {mean_diff:.8f}") + print(f"allclose: {precision_flag}") + + for _ in range(20): + torch_model(*inputs) + cuda_model(*inputs) + + num_iterations = 200 + + torch.cuda.synchronize() + start = time.time() + for _ in range(num_iterations): + torch_model(*inputs) + torch.cuda.synchronize() + torch_time = (time.time() - start) / num_iterations + + torch.cuda.synchronize() + start = time.time() + for _ in range(num_iterations): + cuda_model(*inputs) + torch.cuda.synchronize() + cuda_time = (time.time() - start) / num_iterations + + speedup = torch_time / cuda_time if cuda_time > 0 else 0.0 + print("-------------------- performance check --------------------") + print(f"PyTorch GEGLU average time: {torch_time:.6f} s") + print(f"Custom CUDA GEGLU average time: {cuda_time:.6f} s") + print(f"Speedup: {speedup:.2f}x") + + return precision_flag, speedup + + +if __name__ == "__main__": + run_benchmark() diff --git a/scripts/operator_targets.txt b/scripts/operator_targets.txt index 44e89c9..55a7f01 100644 --- a/scripts/operator_targets.txt +++ b/scripts/operator_targets.txt @@ -195,6 +195,7 @@ gaussian_filter_2d gaussian_pdf GaussianNLLLoss GDL +geglu GELU-Affine-Gate gelu_dropout gempool