InferenceWorker runs continuous batching over a queue of inputs. Continuous batching workflow: 1. Process inputs one at a time from queue 2. Prefill input and insert into KV cache 3. Continue prefilling until enough samples for batch decode 4. Decode until at least one sequ
| 285 | |
| 286 | |
| 287 | class InferenceWorker: |
| 288 | """ |
| 289 | InferenceWorker runs continuous batching over |
| 290 | a queue of inputs. |
| 291 | |
| 292 | Continuous batching workflow: |
| 293 | 1. Process inputs one at a time from queue |
| 294 | 2. Prefill input and insert into KV cache |
| 295 | 3. Continue prefilling until enough samples for batch decode |
| 296 | 4. Decode until at least one sequence completes |
| 297 | 5. Refill newly available decode slots with prefill |
| 298 | 6. Repeat until all sequences complete |
| 299 | |
| 300 | Prefill Packing: |
| 301 | When enable_batch_prefill is True, the prefill processor |
| 302 | will pack multiple inputs into a single sequence before |
| 303 | doing the prefill. |
| 304 | |
| 305 | There are multiple buckets for packed sequences, where each bucket |
| 306 | contains inputs with the same padded length. Only inputs with the same |
| 307 | padded length can be packed together. |
| 308 | |
| 309 | It is important to sort the inputs by padded length so that the |
| 310 | buckets fill up quickly. |
| 311 | |
| 312 | When a decode slot frees up, the prefill processor will add the |
| 313 | sequence to a bucket. If the bucket becomes full, the packed sequence |
| 314 | will be prefilled. |
| 315 | |
| 316 | E.g. |
| 317 | Bucket for length 64: [...seq1, ...seq2, ...seq3, ...seq4] |
| 318 | Bucket for length 128: [...seq1, ...seq2] |
| 319 | Bucket for length 256: [...seq1] |
| 320 | |
| 321 | """ |
| 322 | |
| 323 | def __init__( |
| 324 | self, |
| 325 | config: MaxTextConfig, |
| 326 | params: Params | None, |
| 327 | min_decode_steps: int, |
| 328 | enable_batch_prefill: bool, |
| 329 | devices: list[Any], |
| 330 | tokenizer: Any, |
| 331 | eos_ids: list[int], |
| 332 | prefill_lengths: list[int], |
| 333 | max_decode_length: int, |
| 334 | batch_prefill_max_batch_size: int, |
| 335 | is_pw_reshard: bool = True, |
| 336 | rng: jax.random.PRNGKey = None, |
| 337 | mesh: Mesh = None, |
| 338 | debug: bool = False, |
| 339 | ): |
| 340 | """ |
| 341 | Args: |
| 342 | config: MaxText configuration |
| 343 | params: Model parameters, if None, the params will be loaded from the config |
| 344 | min_decode_steps: Minimum number of decode steps to run at once |