Main function to run interleaved inference.
(argv: Sequence[str])
| 41 | |
| 42 | |
| 43 | def main(argv: Sequence[str]) -> None: |
| 44 | """Main function to run interleaved inference.""" |
| 45 | jax.config.update("jax_default_prng_impl", "unsafe_rbg") |
| 46 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0" |
| 47 | |
| 48 | config = pyconfig.initialize(argv) |
| 49 | _validate_config(config) |
| 50 | max_utils.print_system_information() |
| 51 | |
| 52 | engine = maxengine.MaxEngine(config) |
| 53 | rng = jax.random.PRNGKey(1234) |
| 54 | rng, rng_load_params = jax.random.split(rng) |
| 55 | params = engine.load_params(rng=rng_load_params) |
| 56 | |
| 57 | text = config.prompt |
| 58 | metadata = engine.get_tokenizer() |
| 59 | tokenizer_model = engine.build_tokenizer(metadata) |
| 60 | tokens, true_length = tokenizer_model.encode(text, is_bos=True, prefill_lengths=[config.max_prefill_predict_length]) |
| 61 | assert true_length <= config.max_prefill_predict_length, "Prompt too long for prefill length" |
| 62 | |
| 63 | batch_size = int(config.per_device_batch_size * jax.device_count()) |
| 64 | assert ( |
| 65 | 0 < _NUM_STREAMS <= batch_size |
| 66 | ), f"The number of streams {_NUM_STREAMS} must be > 0 and <= batch size {batch_size}" |
| 67 | |
| 68 | # Initialize decode state |
| 69 | rng, rng_init_decode = jax.random.split(rng) |
| 70 | decode_state = engine.init_decode_state(rng=rng_init_decode) |
| 71 | print("Initial decode state initialized.") |
| 72 | |
| 73 | # Keep track of results per stream (slot) |
| 74 | streams_results: dict[int, list[int]] = {i: [] for i in range(_NUM_STREAMS)} |
| 75 | streams_active: list[bool] = [False] * _NUM_STREAMS # Track which slots are active |
| 76 | streams_finished: list[bool] = [False] * _NUM_STREAMS # Track finished streams |
| 77 | streams_prefilled_count = 0 |
| 78 | streams_inserted_count = 0 |
| 79 | |
| 80 | # --- Initial Prefill Phase --- |
| 81 | print(f"Starting initial prefill for {_INITIAL_PREFILL_STREAMS} streams...") |
| 82 | prefill_results_to_insert = {} # Store prefill results before inserting |
| 83 | for i in range(_INITIAL_PREFILL_STREAMS): |
| 84 | slot_idx = i |
| 85 | print(f" Prefilling stream for slot {slot_idx}...") |
| 86 | rng, rng_prefill = jax.random.split(rng) |
| 87 | request_id = uuid.uuid4() |
| 88 | prefill_result, first_token = engine.prefill( |
| 89 | params=params, |
| 90 | padded_tokens=tokens, |
| 91 | true_length=true_length, |
| 92 | rng=rng_prefill, |
| 93 | slot=slot_idx, |
| 94 | request_id=request_id, |
| 95 | ) |
| 96 | prefill_results_to_insert[slot_idx] = prefill_result |
| 97 | streams_results[slot_idx].append(first_token.get_result_at_slot(0).tokens.item()) |
| 98 | streams_prefilled_count += 1 |
| 99 | print(f"After prefill stream {slot_idx}") |
| 100 |
nothing calls this directly
no test coverage detected