| 30 | ) |
| 31 | ) |
| 32 | def benchmark(T, provider): |
| 33 | device = 'cuda' |
| 34 | dtype = torch.bfloat16 |
| 35 | requires_grad = True |
| 36 | B, H, V = 4, 4096, 120000 |
| 37 | |
| 38 | x = torch.randn(B * T, H, device=device, requires_grad=requires_grad, dtype=dtype) |
| 39 | target = torch.randint(0, V, (B * T,), device=device, dtype=torch.int64) |
| 40 | w = torch.randn(V, H, device=device, requires_grad=requires_grad, dtype=dtype) |
| 41 | b = torch.randn(V, device=device, requires_grad=requires_grad, dtype=dtype) |
| 42 | |
| 43 | quantiles = [0.5, 0.2, 0.8] |
| 44 | results = 0, 0, 0 |
| 45 | if provider == 'naive': |
| 46 | criterion = nn.CrossEntropyLoss() |
| 47 | results = triton.testing.do_bench(lambda: criterion(F.linear(x, w, b), target), quantiles=quantiles) |
| 48 | elif provider == 'naive_bwd': |
| 49 | criterion = nn.CrossEntropyLoss() |
| 50 | results = triton.testing.do_bench(lambda: criterion(F.linear(x, w, b), target).backward(), quantiles=quantiles) |
| 51 | elif provider == 'fused': |
| 52 | criterion = FusedCrossEntropyLoss() |
| 53 | results = triton.testing.do_bench(lambda: criterion(F.linear(x, w, b), target), quantiles=quantiles) |
| 54 | elif provider == 'fused_bwd': |
| 55 | criterion = FusedCrossEntropyLoss() |
| 56 | results = triton.testing.do_bench(lambda: criterion(F.linear(x, w, b), target).backward(), quantiles=quantiles) |
| 57 | elif provider == 'fused_linear': |
| 58 | criterion = FusedLinearCrossEntropyLoss() |
| 59 | results = triton.testing.do_bench(lambda: criterion(x, target, w, b), quantiles=quantiles) |
| 60 | elif provider == 'fused_linear_bwd': |
| 61 | criterion = FusedLinearCrossEntropyLoss() |
| 62 | results = triton.testing.do_bench(lambda: criterion(x, target, w, b).backward(), quantiles=quantiles) |
| 63 | return results |
| 64 | |
| 65 | |
| 66 | if __name__ == '__main__': |