()
| 201 | |
| 202 | |
| 203 | def test_mtp_sampler(): |
| 204 | batch_size = 32 |
| 205 | vocab_size = 1024 |
| 206 | min_seq_len = 1 |
| 207 | max_seq_len = 1024 |
| 208 | max_model_len = 1024 |
| 209 | max_draft_token_num = 1 |
| 210 | |
| 211 | fd_config = _create_fd_config(max_model_len) |
| 212 | sampling_metadata = _create_default_sampling_metadata(batch_size, min_seq_len, max_seq_len) |
| 213 | logits = _create_fake_logits(batch_size * (max_draft_token_num + 1), vocab_size) |
| 214 | |
| 215 | share_inputs = _create_share_inputs(batch_size, max_draft_token_num, max_model_len, vocab_size) |
| 216 | |
| 217 | sampler = MTPSampler(fd_config) |
| 218 | sampler(logits, sampling_metadata, max_model_len, share_inputs) |
| 219 | |
| 220 | |
| 221 | def test_mtp_sampler_logprobs(): |
no test coverage detected