()
| 4 | |
| 5 | |
| 6 | def test_gemv(): |
| 7 | device = torch.device("cuda") |
| 8 | hidden_size = 4096 |
| 9 | vocab_size = 128_000 |
| 10 | hidden_states = torch.randn(hidden_size, device=device) |
| 11 | weights = torch.randn(hidden_size, vocab_size, device=device) |
| 12 | expected = hidden_states @ weights |
| 13 | |
| 14 | tl_result = tl_gemv(hidden_states, weights) |
| 15 | assert torch.cosine_similarity(expected, tl_result, dim=0) > 0.99 |