File size: 674 Bytes
57c3a10 3a9af28 57c3a10 3a9af28 57c3a10 3a9af28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import torch
import flashinfer
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("GELU and multiply operation completed. Output shape:", out.shape)
print("Output tensor sample:", out[:, :5]) # Show first 5 elements
|