Add GEGLU operator sample

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
 
 <details>
 <summary>展开查看算子实现状态</summary>
@@ -231,6 +231,7 @@
 | [x] | gaussian_pdf | [ZZZJ_#120](S1%20codes/ZZZJ_%23120) |
 | [x] | GaussianNLLLoss | [hli28146_#113](S1%20codes/hli28146_%23113)<br>[uucoco 31](S1%20codes/uucoco%2031)<br>[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) |

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/extension.h>
+            torch::Tensor geglu_cuda(torch::Tensor input);
+            """,
+            cuda_sources="""
+            #include <torch/extension.h>
+            #include <cuda_runtime.h>
+
+            __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<<<blocks, threads>>>(
+                    x.data_ptr<float>(),
+                    output.data_ptr<float>(),
+                    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)

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 []

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.

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()

scripts/operator_targets.txt

@@ -195,6 +195,7 @@ gaussian_filter_2d
 gaussian_pdf
 GaussianNLLLoss
 GDL
+geglu
 GELU-Affine-Gate
 gelu_dropout
 gempool