Inner loop for benchmarking prefill step.
(engine_prefill, params, tokens, true_length, iters, num_samples: int | None = None)
| 40 | |
| 41 | |
| 42 | def prefill_benchmark_loop(engine_prefill, params, tokens, true_length, iters, num_samples: int | None = None): |
| 43 | """Inner loop for benchmarking prefill step.""" |
| 44 | start = datetime.datetime.now() |
| 45 | rng = jax.random.PRNGKey(1234) |
| 46 | prefill_result = None |
| 47 | for _ in range(iters): |
| 48 | rng, rng_prefill = jax.random.split(rng) |
| 49 | if num_samples is None: |
| 50 | prefill_result, _ = engine_prefill(params, tokens, true_length, rng_prefill) |
| 51 | else: |
| 52 | prefill_result, _ = engine_prefill[num_samples](params, tokens, true_length, rng_prefill, None) |
| 53 | jax.block_until_ready(prefill_result) |
| 54 | end = datetime.datetime.now() |
| 55 | del prefill_result |
| 56 | return (end - start).total_seconds() |
| 57 | |
| 58 | |
| 59 | def prefill_benchmark(config, engine_prefill, params, tokens, true_length, num_model_params, iters): |
no outgoing calls
no test coverage detected