dags/inference/maxtext_inference.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A DAG to run MaxText inference benchmarks with nightly version."""

import datetime
import numpy as np
from airflow import models
from dags import composer_env
from dags.common.vm_resource import TpuVersion
from dags.inference.maxtext_model_config_generator import generate_model_configs

USER_PREFIX = ""

MAXTEXT_BRANCH = ""
JETSTREAM_BRANCH = ""

maxtext_branch = "" if not MAXTEXT_BRANCH else f"-b {MAXTEXT_BRANCH}"
jetstream_branch = "" if not JETSTREAM_BRANCH else f"-b {JETSTREAM_BRANCH}"

# Run once a day at 8 am UTC (12 am PST)
SCHEDULED_TIME = "0 8 * * *" if composer_env.is_prod_env() else None

LLAMA2_7B = "llama2-7b"
LLAMA2_13B = "llama2-13b"
LLAMA2_70B = "llama2-70b"
GEMMA_7B = "gemma-7b"
MIXTRAL_8_7B = "mixtral-8x7b"

BASE_MODE = "base"
CHAT_MODE = "chat"
INSTRUCT_MODE = "instruct"

W_BF16_KV_BF16 = "w-b16-kv-b16"
W_INT8_KV_INT8 = "w-i8-kv-i8"

CKPT = {
    LLAMA2_7B: {
        BASE_MODE: "gs://inference-benchmarks/models/llama2-7b/2024-04-25-14-01/param-only-decode-ckpt-maxtext/checkpoints/0/items",
        CHAT_MODE: "gs://inference-benchmarks/models/llama2-7b-chat/2024-05-24-12-39/param-only-decode-ckpt-maxtext/checkpoints/0/items",
    },
    LLAMA2_13B: {
        BASE_MODE: "gs://inference-benchmarks/models/llama2-13b/2024-04-25-14-01/param-only-decode-ckpt-maxtext/checkpoints/0/items",
        CHAT_MODE: "gs://inference-benchmarks/models/llama2-13b-chat/2024-05-24-12-39/param-only-decode-ckpt-maxtext/checkpoints/0/items",
    },
    LLAMA2_70B: {
        CHAT_MODE: "gs://inference-benchmarks/models/llama2-70b-chat/2024-05-08-23-16/param-only-decode-ckpt-maxtext/checkpoints/0/items"
    },
    GEMMA_7B: {
        BASE_MODE: "gs://inference-benchmarks/models/gemma-7b/2024-04-25-14-01/param-only-decode-ckpt-maxtext/checkpoints/0/items"
    },
    MIXTRAL_8_7B: {
        # checkpoint created using these instructions - go/mixtral-inference-testing
        INSTRUCT_MODE: "gs://vipannalla_mixtral_ckpt/moe_matmul/moe_matmul_06_15_24/checkpoints/0/items/"
    },
}


dag_id = (
    "jetstream_benchmark_serving"
    if not USER_PREFIX
    else f"{USER_PREFIX}_jetstream_benchmark_serving"
)
tags = ["inference_team", "jetstream", "maxtext", "benchmark"]
if USER_PREFIX:
  tags.append(USER_PREFIX)


with models.DAG(
    dag_id=dag_id,
    tags=tags,
    start_date=datetime.datetime(2024, 1, 19),
    schedule=SCHEDULED_TIME,
    catchup=False,
) as dag:
  test_name_prefix = "max-js" if not USER_PREFIX else f"{USER_PREFIX}-max-js"

  # TODO: baseline layout can be deleted if meeting one of the below conditions:
  #   - default cache layout performs better
  #   - run layout tuning to get optimized cache layout for 0213

  llama2_7B_bf16_batch_sizes = [1, 2, 4, 8, 12]
  llama2_7B_int8_batch_sizes = [1, 2, 4, 8, 12, 24]

  llama2_13B_bf16_batch_sizes = [1, 2, 4, 8]
  llama2_13B_int8_batch_sizes = [1, 2, 4, 8]

  llama2_70B_bf16_batch_sizes = [1, 4, 8, 16, 24]
  llama2_70B_int8_batch_sizes = [1, 2, 4, 8, 11]

  gemma_7B_bf16_batch_sizes = [1, 2, 4, 8, 12]
  gemma_7B_int8_batch_sizes = [1, 2, 4, 8, 12, 24]

  test_templates = {
      # LLAMA2_7B
      LLAMA2_7B: {
          "maxtext_branch": maxtext_branch,
          "jetstream_branch": jetstream_branch,
          "sleep_time": 360,
          "time_out_in_min": 120,
          "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
          "model_name": LLAMA2_7B,
          "tokenizer": "tokenizer.llama2",
          "weight_dtype": "bfloat16",
          "scan_layers": "false",
          "max_prefill_predict_length": 1024,
          "max_target_length": 2048,
          "reshape_q": True,
          # (ici_fsdp_parallelism, ici_autoregressive_parallelism, ici_tensor_parallelism)
          "ici_parallelisms": [(1, 1, -1)],
          "dataset": "openorca",
          "num_prompts": 1000,
          "max_output_length": 1024,
          "warmup_mode": "full",
      },
      f"{LLAMA2_7B}-{W_BF16_KV_BF16}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-2013-2013",  # optimized layout for 0123
              "0213-0213-0213",  # default layout
              "0213-0213-0132",  # optimized layout for 0213
          ],
      },
      f"{LLAMA2_7B}-{W_INT8_KV_INT8}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0213-0213-0213",  # default layout
              "0213-0231-0213",  # optimized layout for 0213
          ],
      },
      # LLAMA2_13B
      LLAMA2_13B: {
          "maxtext_branch": maxtext_branch,
          "jetstream_branch": jetstream_branch,
          "sleep_time": 360,
          "time_out_in_min": 120,
          "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
          "model_name": LLAMA2_13B,
          "tokenizer": "tokenizer.llama2",
          "weight_dtype": "bfloat16",
          "scan_layers": "false",
          "max_prefill_predict_length": 1024,
          "max_target_length": 2048,
          "reshape_q": True,
          # (ici_fsdp_parallelism, ici_autoregressive_parallelism, ici_tensor_parallelism)
          "ici_parallelisms": [(1, 1, -1)],
          "dataset": "openorca",
          "request_rate": [0.0],
          "num_prompts": 1000,
          "max_output_length": 1024,
          "warmup_mode": "full",
      },
      f"{LLAMA2_13B}-{W_BF16_KV_BF16}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      f"{LLAMA2_13B}-{W_INT8_KV_INT8}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      # LLAMA2_70B
      LLAMA2_70B: {
          "maxtext_branch": maxtext_branch,
          "jetstream_branch": jetstream_branch,
          "sleep_time": 360,
          "time_out_in_min": 240,
          "tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)],
          "model_name": LLAMA2_70B,
          "tokenizer": "tokenizer.llama2",
          "weight_dtype": "bfloat16",
          "scan_layers": "false",
          "max_prefill_predict_length": 1024,
          "max_target_length": 2048,
          "reshape_q": True,
          # (ici_fsdp_parallelism, ici_autoregressive_parallelism, ici_tensor_parallelism)
          "ici_parallelisms": [(1, 1, -1)],
          "dataset": "openorca",
          "num_prompts": 1000,
          "max_output_length": 1024,
          "warmup_mode": "full",
      },
      f"{LLAMA2_70B}-{W_BF16_KV_BF16}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      f"{LLAMA2_70B}-{W_INT8_KV_INT8}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      # GEMMA_7B
      GEMMA_7B: {
          "maxtext_branch": maxtext_branch,
          "jetstream_branch": jetstream_branch,
          "sleep_time": 360,
          "time_out_in_min": 120,
          "tpu_version_cores": [(TpuVersion.V5E, 8), (TpuVersion.TRILLIUM, 8)],
          "model_name": GEMMA_7B,
          "tokenizer": "tokenizer.gemma",
          "weight_dtype": "bfloat16",
          "scan_layers": "false",
          "max_prefill_predict_length": 1024,
          "max_target_length": 2048,
          "reshape_q": True,
          # (ici_fsdp_parallelism, ici_autoregressive_parallelism, ici_tensor_parallelism)
          "ici_parallelisms": [(1, 1, -1)],
          "dataset": "sharegpt",
          "dataset_path": "~/ShareGPT_V3_unfiltered_cleaned_split.json",
          "request_rate": [0.0],
          "num_prompts": 1000,
          "max_output_length": 1024,
          "warmup_mode": "full",
      },
      f"{GEMMA_7B}-{W_BF16_KV_BF16}-autoselect": {
          "attention": "autoselected",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      f"{GEMMA_7B}-{W_INT8_KV_INT8}-autoselect": {
          "attention": "autoselected",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      # MIXTRAL_8_7B
      MIXTRAL_8_7B: {
          "maxtext_branch": maxtext_branch,
          "jetstream_branch": jetstream_branch,
          "sleep_time": 240,
          "time_out_in_min": 240,
          "tpu_version_cores": [(TpuVersion.V5P, 8), (TpuVersion.TRILLIUM, 8)],
          "model_name": MIXTRAL_8_7B,
          "tokenizer": "gs://maxtext-external/mixtral-8x7B-v0.1-Instruct/tokenizer.mistral",
          "weight_dtype": "bfloat16",
          "scan_layers": "false",
          "max_prefill_predict_length": 2048,
          "max_target_length": 3072,
          "reshape_q": True,
          # (ici_fsdp_parallelism, ici_autoregressive_parallelism, ici_tensor_parallelism)
          "ici_parallelisms": [(1, 1, -1)],
          "dataset": "openorca",
          "num_prompts": 1000,
          "max_output_length": 1024,
          "warmup_mode": "full",
      },
      f"{MIXTRAL_8_7B}-{W_BF16_KV_BF16}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
      f"{MIXTRAL_8_7B}-{W_INT8_KV_INT8}-dot-product": {
          "attention": "dot_product",
          "request_rate": [0.0],
          "axis_order": [
              "0123-1203-1203",  # baseline layout
              "0213-0213-0213",  # default layout
          ],
      },
  }

  tests_llama2_7b_bf16_base_mode_tests = {}
  tests_llama2_7b_bf16_chat_mode_tests = {}

  tests_llama2_7b_int8_base_mode_tests = {}
  tests_llama2_7b_int8_chat_mode_tests = {}

  # llama 2 7B bfloat16 tests in both base and chat mode
  for bs in llama2_7B_bf16_batch_sizes:
    tests_llama2_7b_bf16_base_mode_tests[
        f"{LLAMA2_7B}-{BASE_MODE}-{W_BF16_KV_BF16}-{bs}"
    ] = (
        test_templates[LLAMA2_7B]
        | test_templates[f"{LLAMA2_7B}-{W_BF16_KV_BF16}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_7B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_BF16_KV_BF16,
            "quantization": "",
            "quantize_kvcache": "false",
            "per_device_batch_size": bs,
            "kv_quant_axis": "",
            "run_eval": False,
        }
    )

    tests_llama2_7b_bf16_chat_mode_tests[
        f"{LLAMA2_7B}-{CHAT_MODE}-{W_BF16_KV_BF16}-{bs}"
    ] = (
        test_templates[LLAMA2_7B]
        | test_templates[f"{LLAMA2_7B}-{W_BF16_KV_BF16}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_7B][CHAT_MODE],
            "model_mode": CHAT_MODE,
            "quant_mode": W_BF16_KV_BF16,
            "quantization": "",
            "quantize_kvcache": "false",
            "per_device_batch_size": bs,
            "kv_quant_axis": "",
            "run_eval": True,
        }
    )

  # llama 2 7B int8 tests in both base and chat mode
  for bs in llama2_7B_int8_batch_sizes:
    tests_llama2_7b_int8_base_mode_tests[
        f"{LLAMA2_7B}-{BASE_MODE}-{W_INT8_KV_INT8}-{bs}"
    ] = (
        test_templates[LLAMA2_7B]
        | test_templates[f"{LLAMA2_7B}-{W_INT8_KV_INT8}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_7B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_INT8_KV_INT8,
            "quantization": "int8",
            "quantize_kvcache": "true",
            "kv_quant_dtype": "int8",
            "per_device_batch_size": bs,
            "kv_quant_axis": "heads_and_dkv",
            "run_eval": False,
        }
    )

    tests_llama2_7b_int8_chat_mode_tests[
        f"{LLAMA2_7B}-{CHAT_MODE}-{W_INT8_KV_INT8}-{bs}"
    ] = (
        test_templates[LLAMA2_7B]
        | test_templates[f"{LLAMA2_7B}-{W_BF16_KV_BF16}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_7B][CHAT_MODE],
            "model_mode": CHAT_MODE,
            "quant_mode": W_INT8_KV_INT8,
            "quantization": "int8",
            "quantize_kvcache": "true",
            "kv_quant_dtype": "int8",
            "per_device_batch_size": bs,
            "kv_quant_axis": "heads_and_dkv",
            "run_eval": True,
        }
    )

  tests_llama2_13b_bf16_base_mode_tests = {}
  tests_llama2_13b_bf16_chat_mode_tests = {}

  tests_llama2_13b_int8_base_mode_tests = {}
  tests_llama2_13b_int8_chat_mode_tests = {}

  # llama 13B 7B bfloat16 tests in both base and chat mode

  for bs in llama2_13B_bf16_batch_sizes:
    tests_llama2_13b_bf16_base_mode_tests[
        f"{LLAMA2_13B}-{BASE_MODE}-{W_BF16_KV_BF16}-{bs}"
    ] = (
        test_templates[LLAMA2_13B]
        | test_templates[f"{LLAMA2_13B}-{W_BF16_KV_BF16}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_13B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_BF16_KV_BF16,
            "quantization": "",
            "quantize_kvcache": "false",
            "per_device_batch_size": bs,
            "kv_quant_axis": "",
            "run_eval": False,
        }
    )

    tests_llama2_13b_bf16_chat_mode_tests[
        f"{LLAMA2_13B}-{BASE_MODE}-{W_BF16_KV_BF16}-{bs}"
    ] = (
        test_templates[LLAMA2_13B]
        | test_templates[f"{LLAMA2_13B}-{W_BF16_KV_BF16}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_13B][CHAT_MODE],
            "model_mode": CHAT_MODE,
            "quant_mode": W_BF16_KV_BF16,
            "quantization": "",
            "quantize_kvcache": "false",
            "per_device_batch_size": bs,
            "kv_quant_axis": "",
            "run_eval": True,
        }
    )

  for bs in llama2_13B_int8_batch_sizes:
    tests_llama2_13b_int8_base_mode_tests[
        f"{LLAMA2_13B}-{BASE_MODE}-{W_INT8_KV_INT8}-{bs}"
    ] = (
        test_templates[LLAMA2_13B]
        | test_templates[f"{LLAMA2_13B}-{W_INT8_KV_INT8}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_13B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_INT8_KV_INT8,
            "quantization": "int8",
            "quantize_kvcache": "true",
            "kv_quant_dtype": "int8",
            "per_device_batch_size": bs,
            "kv_quant_axis": "heads_and_dkv",
            "run_eval": False,
        }
    )

    tests_llama2_13b_int8_chat_mode_tests[
        f"{LLAMA2_13B}-{CHAT_MODE}-{W_INT8_KV_INT8}-{bs}"
    ] = (
        test_templates[LLAMA2_13B]
        | test_templates[f"{LLAMA2_13B}-{W_INT8_KV_INT8}-dot-product"]
        | {
            "checkpoint": CKPT[LLAMA2_13B][CHAT_MODE],
            "model_mode": CHAT_MODE,
            "quant_mode": W_INT8_KV_INT8,
            "quantization": "int8",
            "quantize_kvcache": "true",
            "kv_quant_dtype": "int8",
            "per_device_batch_size": bs,
            "kv_quant_axis": "heads_and_dkv",
            "run_eval": True,
        }
    )

  tests_gemma_7b_bf16_base_mode_tests = {}
  tests_gemma_7b_int8_base_mode_tests = {}

  for bs in gemma_7B_bf16_batch_sizes:
    tests_gemma_7b_bf16_base_mode_tests[
        f"{GEMMA_7B}-{BASE_MODE}-{W_BF16_KV_BF16}-{bs}"
    ] = (
        test_templates[GEMMA_7B]
        | test_templates[f"{GEMMA_7B}-{W_BF16_KV_BF16}-autoselect"]
        | {
            "checkpoint": CKPT[GEMMA_7B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_BF16_KV_BF16,
            "quantization": "",
            "quantize_kvcache": "false",
            "per_device_batch_size": bs,
            "kv_quant_axis": "",
            "run_eval": False,
        }
    )

  for bs in gemma_7B_int8_batch_sizes:
    tests_gemma_7b_int8_base_mode_tests[
        f"{GEMMA_7B}-{BASE_MODE}-{W_INT8_KV_INT8}-{bs}"
    ] = (
        test_templates[GEMMA_7B]
        | test_templates[f"{GEMMA_7B}-{W_INT8_KV_INT8}-autoselect"]
        | {
            "checkpoint": CKPT[GEMMA_7B][BASE_MODE],
            "model_mode": BASE_MODE,
            "quant_mode": W_INT8_KV_INT8,
            "quantization": "int8",
            "quantize_kvcache": "true",
            "kv_quant_dtype": "int8",
            "per_device_batch_size": bs,
            "kv_quant_axis": "heads_and_dkv",
            "run_eval": False,
        }
    )

  tests = (
      tests_llama2_7b_bf16_base_mode_tests
      | tests_llama2_7b_bf16_chat_mode_tests
      | tests_llama2_7b_int8_base_mode_tests
      | tests_llama2_7b_int8_chat_mode_tests
      | tests_llama2_13b_bf16_base_mode_tests
      | tests_llama2_13b_bf16_chat_mode_tests
      | tests_llama2_13b_int8_base_mode_tests
      | tests_llama2_13b_int8_chat_mode_tests
      | tests_gemma_7b_bf16_base_mode_tests
      | tests_gemma_7b_int8_base_mode_tests
      | {
          # LLAMA2_70B
          f"{LLAMA2_70B}-{CHAT_MODE}-{W_BF16_KV_BF16}": test_templates[
              LLAMA2_70B
          ]
          | test_templates[f"{LLAMA2_70B}-{W_BF16_KV_BF16}-dot-product"]
          | {
              "checkpoint": CKPT[LLAMA2_70B][CHAT_MODE],
              "model_mode": CHAT_MODE,
              "quant_mode": W_BF16_KV_BF16,
              "quantization": "",
              "quantize_kvcache": "false",
              "per_device_batch_size": 24,
              "kv_quant_axis": "",
              "run_eval": True,
          },
          f"{LLAMA2_70B}-{CHAT_MODE}-{W_INT8_KV_INT8}": test_templates[
              LLAMA2_70B
          ]
          | test_templates[f"{LLAMA2_70B}-{W_INT8_KV_INT8}-dot-product"]
          | {
              "checkpoint": CKPT[LLAMA2_70B][CHAT_MODE],
              "model_mode": CHAT_MODE,
              "quant_mode": W_INT8_KV_INT8,
              "quantization": "int8",
              "quantize_kvcache": "true",
              "kv_quant_dtype": "int8",
              "per_device_batch_size": 48,
              "kv_quant_axis": "heads_and_dkv",
              "run_eval": True,
          },
          # MIXTRAL_8_7B
          f"{MIXTRAL_8_7B}-{INSTRUCT_MODE}-{W_BF16_KV_BF16}": test_templates[
              MIXTRAL_8_7B
          ]
          | test_templates[f"{MIXTRAL_8_7B}-{W_BF16_KV_BF16}-dot-product"]
          | {
              "checkpoint": CKPT[MIXTRAL_8_7B][INSTRUCT_MODE],
              "model_mode": INSTRUCT_MODE,
              "quant_mode": W_BF16_KV_BF16,
              "quantization": "",
              "quantize_kvcache": "false",
              "per_device_batch_size": 128,
              "kv_quant_axis": "",
              "run_eval": True,
          },
          f"{MIXTRAL_8_7B}-{INSTRUCT_MODE}-{W_INT8_KV_INT8}": test_templates[
              MIXTRAL_8_7B
          ]
          | test_templates[f"{MIXTRAL_8_7B}-{W_INT8_KV_INT8}-dot-product"]
          | {
              "checkpoint": CKPT[MIXTRAL_8_7B][INSTRUCT_MODE],
              "model_mode": INSTRUCT_MODE,
              "quant_mode": W_INT8_KV_INT8,
              "quantization": "int8",
              "quantize_kvcache": "true",
              "kv_quant_dtype": "int8",
              "per_device_batch_size": 258,
              "kv_quant_axis": "heads_and_dkv",
              "run_eval": True,
          },
      }
  )

  # run_configs = [
  #     f"{LLAMA2_7B}-{BASE_MODE}-{W_BF16_KV_BF16}",
  #     f"{LLAMA2_7B}-{BASE_MODE}-{W_INT8_KV_INT8}",
  #     f"{LLAMA2_7B}-{CHAT_MODE}-{W_BF16_KV_BF16}",
  #     f"{LLAMA2_7B}-{CHAT_MODE}-{W_INT8_KV_INT8}",
  #     f"{LLAMA2_13B}-{BASE_MODE}-{W_BF16_KV_BF16}",
  #     f"{LLAMA2_13B}-{BASE_MODE}-{W_INT8_KV_INT8}",
  #     f"{LLAMA2_13B}-{CHAT_MODE}-{W_BF16_KV_BF16}",
  #     f"{LLAMA2_13B}-{CHAT_MODE}-{W_INT8_KV_INT8}",
  #     f"{LLAMA2_70B}-{CHAT_MODE}-{W_BF16_KV_BF16}",
  #     f"{LLAMA2_70B}-{CHAT_MODE}-{W_INT8_KV_INT8}",
  #     f"{GEMMA_7B}-{BASE_MODE}-{W_BF16_KV_BF16}",
  #     f"{GEMMA_7B}-{BASE_MODE}-{W_INT8_KV_INT8}",
  #     f"{MIXTRAL_8_7B}-{INSTRUCT_MODE}-{W_BF16_KV_BF16}",
  #     f"{MIXTRAL_8_7B}-{INSTRUCT_MODE}-{W_INT8_KV_INT8}",
  # ]

  skip_configs = []

  dags = []
  for model_config_name, sweep_model_configs in tests.items():
    # if run_configs and model_config_name not in run_configs:
    #   continue
    if skip_configs and model_config_name in skip_configs:
      continue
    for tpu_version, tpu_cores in sweep_model_configs["tpu_version_cores"]:
      for axis_order in sweep_model_configs["axis_order"]:
        for ici_parallelism in sweep_model_configs["ici_parallelisms"]:
          for request_rate in sweep_model_configs["request_rate"]:
            jetstream_benchmark_serving_kv_cache_layout = (
                generate_model_configs(
                    test_name_prefix=test_name_prefix,
                    model_config_name=model_config_name,
                    sweep_model_configs=sweep_model_configs,
                    axis_order=axis_order,
                    ici_parallelism=ici_parallelism,
                    request_rate=request_rate,
                    tpu_version=tpu_version,
                    tpu_cores=tpu_cores,
                )
            )
            dags.append(jetstream_benchmark_serving_kv_cache_layout)

  # Cap the number of simultaneously requested v5e-8 due to resource contraints
  n_parallel_jobs = 4
  chunks = np.array_split(dags, n_parallel_jobs)
  for chunk in chunks:
    for i in range(1, len(chunk)):
      chunk[i - 1] >> chunk[i]