Example inference workload

[rwitten]

python MaxText/decode.py MaxText/configs/base.yml per_device_batch_size=64 run_name=runner_2024-01-30-20-02 max_prefill_predict_length=128 max_target_length=256 dataset_path=gs://maxtext-dataset async_checkpointing=false scan_layers=false attention=dot_product scan_layers=false ici_autoregressive_parallelism=4

400GB/s/device on a v4-8