Run microbenchmarks.
(config)
| 298 | |
| 299 | |
| 300 | def run_benchmarks(config): |
| 301 | """Run microbenchmarks.""" |
| 302 | engine = maxengine.MaxEngine(config) |
| 303 | prefill_processor = prefill_packing.PrefillProcessor(engine) |
| 304 | rng = jax.random.PRNGKey(1234) |
| 305 | rng, rng_load_params = jax.random.split(rng) |
| 306 | params = engine.load_params(rng_load_params) |
| 307 | prefill_lengths = [int(l) for l in config.inference_microbenchmark_prefill_lengths.split(",")] |
| 308 | stages_to_benchmark = config.inference_microbenchmark_stages.split(",") |
| 309 | benchmark_loop_iters = config.inference_microbenchmark_loop_iters |
| 310 | |
| 311 | text = config.prompt |
| 312 | metadata = engine.get_tokenizer() |
| 313 | tokenizer_model = engine.build_tokenizer(metadata) |
| 314 | rng, rng_init_decode = jax.random.split(rng) |
| 315 | |
| 316 | generate_executable, params, decode_state_executable = engine.aot_compile(params, pass_rng_shape=True) |
| 317 | decode_state = decode_state_executable(rng_init_decode) |
| 318 | |
| 319 | _, cache_size, _ = max_utils.summarize_pytree_data(decode_state["cache"], name="Cache") |
| 320 | num_model_params, model_size, _ = max_utils.summarize_pytree_data(params, name="Model") |
| 321 | |
| 322 | benchmark_results = {} |
| 323 | if "prefill" in stages_to_benchmark: |
| 324 | benchmark_results["prefill-result-sizes"] = {} |
| 325 | benchmark_results["prefill"] = {} |
| 326 | benchmark_results["insert"] = {} |
| 327 | prefill_tokens = {} |
| 328 | prefill_true_lengths = {} |
| 329 | prefill_executable = {} |
| 330 | prefill_insert_executable = {} |
| 331 | i32_scalar = jax.ShapeDtypeStruct((), int) |
| 332 | rng_shape = jax.ShapeDtypeStruct([4], jax.numpy.dtype("uint32")) |
| 333 | |
| 334 | for prefill_length in prefill_lengths: |
| 335 | prefill_tokens[prefill_length], prefill_true_lengths[prefill_length] = tokenizer_model.encode( |
| 336 | text, is_bos=True, prefill_lengths=[prefill_length] |
| 337 | ) |
| 338 | |
| 339 | key_shape = jax.ShapeDtypeStruct([prefill_length], jax.numpy.dtype("int32")) |
| 340 | prefill_executable[prefill_length] = ( |
| 341 | jax.jit( |
| 342 | engine.prefill_aot, |
| 343 | in_shardings=(engine.param_layouts, None, None, None), |
| 344 | ).lower(params, key_shape, i32_scalar, rng_shape) |
| 345 | ).compile(compiler_options=None) |
| 346 | |
| 347 | prefill_insert_executable[prefill_length] = prefill_processor.aot_compile(params, prefill_length) |
| 348 | |
| 349 | benchmark_results["prefill-result-sizes"][prefill_length] = summarize_prefill_result( |
| 350 | prefill_executable[prefill_length], params, prefill_tokens[prefill_length], prefill_true_lengths[prefill_length] |
| 351 | ) |
| 352 | |
| 353 | for prefill_length in prefill_lengths: |
| 354 | benchmark_results["prefill"][prefill_length] = prefill_benchmark( |
| 355 | config, |
| 356 | prefill_executable[prefill_length], |
| 357 | params, |