Handles warmup, running prefill benchmark, and printing results.
(config, engine_prefill, params, tokens, true_length, num_model_params, iters)
| 57 | |
| 58 | |
| 59 | def prefill_benchmark(config, engine_prefill, params, tokens, true_length, num_model_params, iters): |
| 60 | """Handles warmup, running prefill benchmark, and printing results.""" |
| 61 | rng = jax.random.PRNGKey(1234) |
| 62 | prefill_result = None |
| 63 | for _ in range(_WARMUP_ITERS): |
| 64 | rng, rng_prefill = jax.random.split(rng) |
| 65 | prefill_result, _ = engine_prefill(params, tokens, true_length, rng_prefill) |
| 66 | jax.block_until_ready(prefill_result) |
| 67 | del prefill_result |
| 68 | |
| 69 | print(f"Prefill benchmark results for length {tokens.size}:\n") |
| 70 | time_in_s = prefill_benchmark_loop(engine_prefill, params, tokens, true_length, iters) |
| 71 | prefill_average_ms = 1000 * time_in_s / iters |
| 72 | prefill_tflops_per_device, _, _ = maxtext_utils.calculate_prefill_tflops_per_device( |
| 73 | num_model_params, tokens.size, config |
| 74 | ) |
| 75 | tflops_per_sec_per_device = prefill_tflops_per_device / prefill_average_ms * 1000.0 |
| 76 | print( |
| 77 | f"\tPrefill step average time: {prefill_average_ms:.3f} ms\n" |
| 78 | f"\tPrefill total TFLOPs/device: {prefill_tflops_per_device:.3f}\n" |
| 79 | f"\tPrefill TFLOPs/sec/device: {tflops_per_sec_per_device:.3f}\n\n\n\n" |
| 80 | ) |
| 81 | result_dict = { |
| 82 | "time_in_ms": prefill_average_ms, |
| 83 | "total_tflops_per_device": prefill_tflops_per_device, |
| 84 | "tflops_per_sec_per_device": tflops_per_sec_per_device, |
| 85 | } |
| 86 | return result_dict |
| 87 | |
| 88 | |
| 89 | def prefill_multisampling_benchmark(config, engine_prefill_multisampling, params, tokens, true_length, iters): |
no test coverage detected