feat: add gelu example

README.md

@@ -7,4 +7,11 @@ tags:
 
 This kernel is a work in progress and requires more work to correctly add all of the FlashInfer kernels.
 
-Please see the [generate-source.md](generate-source.md) for instructions on how to generate the source files that are contained in this kernel.
+Please see the [generate-source.md](generate-source.md) for instructions on how to generate the source files that are contained in this kernel.
+
+
+### Run `gelu_and_mul` operation
+
+```bash
+nix develop -L .#test --command python tests/simple_test.py
+```

tests/simple_test.py

@@ -1,9 +1,21 @@
 import torch
 import flashinfer
 
-out = torch.empty(3, 3, dtype=torch.float16).cuda()
-input_tensor = torch.randn(3, 3, dtype=torch.float16).cuda()
+torch.manual_seed(42)
+
+batch_size = 128
+hidden_dim = 4096
+out = torch.empty(batch_size, hidden_dim // 2, dtype=torch.float16)
+input_tensor = torch.randn(batch_size, hidden_dim, dtype=torch.float16)
 enable_pdl = True
 
+# make sure the output tensor is on the GPU if available
+if torch.cuda.is_available():
+    out = out.cuda()
+    input_tensor = input_tensor.cuda()
+else:
+    raise RuntimeError("CUDA is not available. Please run this on a GPU.")
+
 flashinfer.gelu_and_mul(out, input_tensor, enable_pdl)
-print(out)
+print("GELU and multiply operation completed. Output shape:", out.shape)
+print("Output tensor sample:", out[:, :5])  # Show first 5 elements