Почему реализация LayerNorm в pytorch не приближается к пиковой пропускной способности памяти A100? ⇐ Python

[Anonymous]

Я пытался протестировать torch.nn.LayerNorm на A100, чтобы сравнить его с созданным мной собственным ядром CUDA. Я ожидал, что ядро ​​PyTorch приблизится к пиковой пропускной способности памяти A100 (которая в Интернете показывает, что она составляет около 2 ТБ/с), но цифры, которые я вижу, намного ниже. Я неправильно оцениваю LayerNorm, или есть что-то еще, что мешает реализации пирторча достичь пика?
Это код, который я использовал:
import torch
import torch.utils.benchmark as benchmark

B, H = 8192, 4096
# Use the same kind of data as my CUDA test
x = torch.randn(B, H, device='cuda', dtype=torch.float32)
gamma = torch.ones(H, device='cuda')
beta = torch.zeros(H, device='cuda')
ln = torch.nn.LayerNorm(H, elementwise_affine=True).cuda()

# Match gamma/beta to my test
ln.weight.data = gamma
ln.bias.data = beta

# Warmup
for _ in range(20):
_ = ln(x)
torch.cuda.synchronize()

# Benchmark
t = benchmark.Timer(
stmt='ln(x)',
globals={'ln': ln, 'x': x}
)
result = t.blocked_autorange(min_run_time=1.0)
print(result)

kernel_time_s = result.mean

# Naive bytes moved: read x, write y (float32)
bytes_moved = 2 * B * H * 4
bandwidth_gbs = (bytes_moved / kernel_time_s) / 1e9
print(f"PyTorch bandwidth: {bandwidth_gbs:.1f} GB/s")

# For comparison, my custom kernel gets ~964 GB/s on the same shape/GPU
print(f\"My kernel (separate benchmark): 964.0 GB/s\")

При запуске я вижу
ln(x)
Median: 236.94 us
IQR: 0.31 us (236.80 to 237.11)
5 measurements, 1000 runs per measurement, 1 thread
PyTorch bandwidth: 1132.8 GB/s
Your kernel: 964.0 GB/s


Подробнее здесь: https://stackoverflow.com/questions/798 ... mory-bandw