forked from ccf-ai-infra/GPUCodeForces
16 lines
546 B
Plaintext
16 lines
546 B
Plaintext
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.
|