Verify that greedy_sampling=True returns the argmax token for each sequence.
(vocab_size, n_hidden_states)
| 218 | @pytest.mark.parametrize("n_hidden_states", [1, 2]) |
| 219 | @pytest.mark.parametrize("vocab_size", [100, 200, 256]) |
| 220 | def test_greedy_sampling(vocab_size, n_hidden_states): |
| 221 | """Verify that greedy_sampling=True returns the argmax token for each sequence.""" |
| 222 | inputs = make_synthetic_inputs(vocab_size=vocab_size, n_hidden_states=n_hidden_states) |
| 223 | |
| 224 | samples = fused_mm_sample_triton( |
| 225 | weights=inputs.weights, |
| 226 | hidden_states=inputs.hidden_states, |
| 227 | num_samples=1, |
| 228 | temperature=torch.empty((), device=device), |
| 229 | seed=0, |
| 230 | greedy_sampling=True, |
| 231 | ) |
| 232 | |
| 233 | ref_logits = inputs.hidden_states.float() @ inputs.weights.float().T # [H, V] |
| 234 | expected = ref_logits.argmax(dim=-1) # [H] |
| 235 | assert samples.shape == (n_hidden_states, 1) |
| 236 | torch.testing.assert_close(samples[:, 0], expected) |
| 237 | |
| 238 | |
| 239 | def test_speed_test_smoke(): |
nothing calls this directly
no test coverage detected