tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py

@@ -0,0 +1,609 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tuned block sizes for quantized matmul kernel."""
+
+import re
+from typing import NamedTuple
+
+import jax
+
+from tpu_inference.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class TunedKey(NamedTuple):
+    tpu_version: int
+    n_batch: int
+    n_out: int
+    n_in: int
+    x_q_dtype: str
+    w_q_dtype: str
+
+
+class TunedValue(NamedTuple):
+    batch_block_size: int
+    out_block_size: int
+    in_block_size: int
+
+
+TUNED_BLOCK_SIZES_RAW = {
+    # go/keep-sorted start
+    (6, 1024, 1024, 4096, 'int8', 'int8'): (1024, 256, 4096),
+    (6, 1024, 1024, 8192, 'int8', 'int8'): (1024, 512, 8192),
+    (6, 1024, 128, 8192, 'int8', 'int8'): (512, 128, 8192),
+    (6, 1024, 1280, 8192, 'int8', 'int8'): (1024, 1280, 2048),
+    (6, 1024, 13824, 5120, 'int8', 'int8'): (1024, 768, 5120),
+    (6, 1024, 14336, 4096, 'int8', 'int8'): (1024, 1024, 4096),
+    (6, 1024, 1536, 1536, 'int8', 'int8'): (256, 1536, 1536),
+    (6, 1024, 1536, 8960, 'int8', 'int8'): (1024, 256, 8960),
+    (6, 1024, 1792, 5120, 'int8', 'int8'): (1024, 256, 5120),
+    (6, 1024, 17920, 1536, 'int8', 'int8'): (1024, 1792, 1536),
+    (6, 1024, 2048, 1536, 'int8', 'int8'): (256, 2048, 1536),
+    (6, 1024, 28672, 4096, 'int8', 'int8'): (1024, 2048, 4096),
+    (6, 1024, 3584, 18944, 'int8', 'int8'): (512, 512, 18944),
+    (6, 1024, 3584, 3584, 'int8', 'int8'): (1024, 896, 3584),
+    (6, 1024, 3584, 8192, 'int8', 'int8'): (1024, 512, 8192),
+    (6, 1024, 37888, 3584, 'int8', 'int8'): (1024, 1024, 3584),
+    (6, 1024, 4096, 14336, 'int8', 'int8'): (1024, 256, 14336),
+    (6, 1024, 4096, 4096, 'int8', 'int8'): (1024, 256, 4096),
+    (6, 1024, 4608, 3584, 'int8', 'int8'): (1024, 768, 3584),
+    (6, 1024, 5120, 1280, 'int8', 'int8'): (1024, 1280, 1280),
+    (6, 1024, 5120, 3456, 'int8', 'int8'): (1024, 512, 3456),
+    (6, 1024, 5120, 640, 'int8', 'int8'): (256, 5120, 640),
+    (6, 1024, 5120, 6912, 'int8', 'int8'): (1024, 512, 6912),
+    (6, 1024, 6144, 4096, 'int8', 'int8'): (1024, 768, 4096),
+    (6, 1024, 6912, 5120, 'int8', 'int8'): (1024, 768, 5120),
+    (6, 1024, 7168, 8192, 'int8', 'int8'): (1024, 256, 8192),
+    (6, 1024, 8192, 1024, 'int8', 'int8'): (1024, 2048, 1024),
+    (6, 1024, 8192, 3584, 'int8', 'int8'): (1024, 1024, 3584),
+    (6, 1024, 896, 5120, 'int8', 'int8'): (1024, 896, 2560),
+    (6, 128, 1024, 4096, 'int8', 'int8'): (128, 512, 4096),
+    (6, 128, 1024, 8192, 'int8', 'int8'): (128, 1024, 2048),
+    (6, 128, 128, 8192, 'int8', 'int8'): (128, 128, 4096),
+    (6, 128, 1280, 8192, 'int8', 'int8'): (128, 1280, 2048),
+    (6, 128, 13824, 5120, 'int8', 'int8'): (128, 13824, 256),
+    (6, 128, 14336, 4096, 'int8', 'int8'): (128, 896, 4096),
+    (6, 128, 1536, 1536, 'int8', 'int8'): (128, 768, 1536),
+    (6, 128, 1536, 8960, 'int8', 'int8'): (128, 384, 8960),
+    (6, 128, 1792, 5120, 'int8', 'int8'): (128, 1792, 1280),
+    (6, 128, 17920, 1536, 'int8', 'int8'): (128, 8960, 384),
+    (6, 128, 2048, 1536, 'int8', 'int8'): (128, 1024, 1536),
+    (6, 128, 28672, 4096, 'int8', 'int8'): (128, 28672, 256),
+    (6, 128, 3584, 18944, 'int8', 'int8'): (128, 512, 9472),
+    (6, 128, 3584, 3584, 'int8', 'int8'): (128, 3584, 896),
+    (6, 128, 3584, 8192, 'int8', 'int8'): (128, 3584, 1024),
+    (6, 128, 37888, 3584, 'int8', 'int8'): (128, 1024, 3584),
+    (6, 128, 4096, 14336, 'int8', 'int8'): (128, 256, 14336),
+    (6, 128, 4096, 4096, 'int8', 'int8'): (128, 512, 4096),
+    (6, 128, 4608, 3584, 'int8', 'int8'): (128, 768, 3584),
+    (6, 128, 5120, 1280, 'int8', 'int8'): (128, 2560, 1280),
+    (6, 128, 5120, 3456, 'int8', 'int8'): (128, 640, 3456),
+    (6, 128, 5120, 640, 'int8', 'int8'): (128, 5120, 128),
+    (6, 128, 5120, 6912, 'int8', 'int8'): (128, 512, 6912),
+    (6, 128, 6144, 4096, 'int8', 'int8'): (128, 1024, 4096),
+    (6, 128, 6912, 5120, 'int8', 'int8'): (128, 768, 5120),
+    (6, 128, 7168, 8192, 'int8', 'int8'): (128, 3584, 1024),
+    (6, 128, 8192, 1024, 'int8', 'int8'): (128, 4096, 1024),
+    (6, 128, 8192, 3584, 'int8', 'int8'): (128, 8192, 512),
+    (6, 128, 896, 5120, 'int8', 'int8'): (128, 896, 2560),
+    (6, 16, 1024, 4096, 'int8', 'int8'): (16, 1024, 2048),
+    (6, 16, 1024, 8192, 'int8', 'int8'): (16, 256, 8192),
+    (6, 16, 128, 8192, 'int8', 'int8'): (16, 128, 8192),
+    (6, 16, 1280, 8192, 'int8', 'int8'): (16, 256, 8192),
+    (6, 16, 13824, 5120, 'int8', 'int8'): (16, 512, 5120),
+    (6, 16, 14336, 4096, 'int8', 'int8'): (16, 896, 4096),
+    (6, 16, 1536, 1536, 'int8', 'int8'): (16, 768, 1536),
+    (6, 16, 1536, 8960, 'int8', 'int8'): (16, 1536, 1792),
+    (6, 16, 1792, 5120, 'int8', 'int8'): (16, 1792, 2560),
+    (6, 16, 17920, 1536, 'int8', 'int8'): (16, 1792, 1536),
+    (6, 16, 2048, 1536, 'int8', 'int8'): (16, 1024, 1536),
+    (6, 16, 28672, 4096, 'int8', 'int8'): (16, 1024, 4096),
+    (6, 16, 3584, 18944, 'int8', 'int8'): (16, 256, 18944),
+    (6, 16, 3584, 3584, 'int8', 'int8'): (16, 896, 3584),
+    (6, 16, 3584, 8192, 'int8', 'int8'): (16, 3584, 1024),
+    (6, 16, 37888, 3584, 'int8', 'int8'): (16, 1024, 3584),
+    (6, 16, 4096, 14336, 'int8', 'int8'): (16, 2048, 1792),
+    (6, 16, 4096, 4096, 'int8', 'int8'): (16, 1024, 4096),
+    (6, 16, 4608, 3584, 'int8', 'int8'): (16, 1152, 3584),
+    (6, 16, 5120, 1280, 'int8', 'int8'): (16, 2560, 1280),
+    (6, 16, 5120, 3456, 'int8', 'int8'): (16, 1024, 3456),
+    (6, 16, 5120, 640, 'int8', 'int8'): (16, 2560, 640),
+    (6, 16, 5120, 6912, 'int8', 'int8'): (16, 640, 6912),
+    (6, 16, 6144, 4096, 'int8', 'int8'): (16, 768, 4096),
+    (6, 16, 6912, 5120, 'int8', 'int8'): (16, 768, 5120),
+    (6, 16, 7168, 8192, 'int8', 'int8'): (16, 512, 8192),
+    (6, 16, 8192, 1024, 'int8', 'int8'): (16, 2048, 1024),
+    (6, 16, 8192, 3584, 'int8', 'int8'): (16, 1024, 3584),
+    (6, 16, 896, 5120, 'int8', 'int8'): (16, 896, 2560),
+    (6, 16384, 1024, 4096, 'int8', 'int8'): (2048, 1024, 4096),
+    (6, 16384, 1024, 8192, 'int8', 'int8'): (1024, 512, 8192),
+    (6, 16384, 128, 8192, 'int8', 'int8'): (256, 128, 8192),
+    (6, 16384, 1280, 8192, 'int8', 'int8'): (512, 1280, 8192),
+    (6, 16384, 13824, 5120, 'int8', 'int8'): (256, 2304, 5120),
+    (6, 16384, 14336, 4096, 'int8', 'int8'): (1024, 2048, 4096),
+    (6, 16384, 1536, 1536, 'int8', 'int8'): (1024, 1536, 1536),
+    (6, 16384, 1536, 8960, 'int8', 'int8'): (1024, 1536, 8960),
+    (6, 16384, 1792, 5120, 'int8', 'int8'): (512, 1792, 5120),
+    (6, 16384, 17920, 1536, 'int8', 'int8'): (512, 17920, 1536),
+    (6, 16384, 2048, 1536, 'int8', 'int8'): (4096, 2048, 1536),
+    (6, 16384, 28672, 4096, 'int8', 'int8'): (2048, 1792, 4096),
+    (6, 16384, 3584, 18944, 'int8', 'int8'): (2048, 3584, 512),
+    (6, 16384, 3584, 3584, 'int8', 'int8'): (1024, 3584, 3584),
+    (6, 16384, 3584, 8192, 'int8', 'int8'): (1024, 1792, 8192),
+    (6, 16384, 37888, 3584, 'int8', 'int8'): (2048, 1024, 3584),
+    (6, 16384, 4096, 14336, 'int8', 'int8'): (2048, 4096, 512),
+    (6, 16384, 4096, 4096, 'int8', 'int8'): (1024, 4096, 4096),
+    (6, 16384, 4608, 3584, 'int8', 'int8'): (1024, 2304, 3584),
+    (6, 16384, 5120, 1280, 'int8', 'int8'): (1024, 5120, 1280),
+    (6, 16384, 5120, 3456, 'int8', 'int8'): (512, 5120, 3456),
+    (6, 16384, 5120, 640, 'int8', 'int8'): (1024, 5120, 640),
+    (6, 16384, 5120, 6912, 'int8', 'int8'): (512, 5120, 6912),
+    (6, 16384, 6144, 4096, 'int8', 'int8'): (1024, 6144, 4096),
+    (6, 16384, 6912, 5120, 'int8', 'int8'): (512, 6912, 5120),
+    (6, 16384, 7168, 8192, 'int8', 'int8'): (1024, 1792, 4096),
+    (6, 16384, 8192, 1024, 'int8', 'int8'): (256, 8192, 1024),
+    (6, 16384, 8192, 3584, 'int8', 'int8'): (1024, 8192, 3584),
+    (6, 16384, 896, 5120, 'int8', 'int8'): (1024, 896, 2560),
+    (6, 2048, 1024, 4096, 'int8', 'int8'): (2048, 256, 4096),
+    (6, 2048, 1024, 8192, 'int8', 'int8'): (2048, 256, 8192),
+    (6, 2048, 128, 8192, 'int8', 'int8'): (1024, 128, 8192),
+    (6, 2048, 1280, 8192, 'int8', 'int8'): (2048, 256, 8192),
+    (6, 2048, 13824, 5120, 'int8', 'int8'): (1024, 768, 5120),
+    (6, 2048, 14336, 4096, 'int8', 'int8'): (2048, 1024, 4096),
+    (6, 2048, 1536, 1536, 'int8', 'int8'): (2048, 512, 1536),
+    (6, 2048, 1536, 8960, 'int8', 'int8'): (1024, 768, 8960),
+    (6, 2048, 1792, 5120, 'int8', 'int8'): (2048, 256, 5120),
+    (6, 2048, 17920, 1536, 'int8', 'int8'): (2048, 2560, 1536),
+    (6, 2048, 2048, 1536, 'int8', 'int8'): (256, 2048, 1536),
+    (6, 2048, 28672, 4096, 'int8', 'int8'): (2048, 1024, 4096),
+    (6, 2048, 3584, 18944, 'int8', 'int8'): (2048, 3584, 512),
+    (6, 2048, 3584, 3584, 'int8', 'int8'): (2048, 512, 3584),
+    (6, 2048, 3584, 8192, 'int8', 'int8'): (2048, 512, 8192),
+    (6, 2048, 37888, 3584, 'int8', 'int8'): (2048, 1024, 3584),
+    (6, 2048, 4096, 14336, 'int8', 'int8'): (2048, 4096, 512),
+    (6, 2048, 4096, 4096, 'int8', 'int8'): (2048, 512, 4096),
+    (6, 2048, 4608, 3584, 'int8', 'int8'): (2048, 512, 3584),
+    (6, 2048, 5120, 1280, 'int8', 'int8'): (2048, 1280, 1280),
+    (6, 2048, 5120, 3456, 'int8', 'int8'): (2048, 512, 3456),
+    (6, 2048, 5120, 640, 'int8', 'int8'): (256, 5120, 640),
+    (6, 2048, 5120, 6912, 'int8', 'int8'): (2048, 512, 6912),
+    (6, 2048, 6144, 4096, 'int8', 'int8'): (2048, 512, 4096),
+    (6, 2048, 6912, 5120, 'int8', 'int8'): (2048, 768, 5120),
+    (6, 2048, 7168, 8192, 'int8', 'int8'): (1024, 512, 8192),
+    (6, 2048, 8192, 1024, 'int8', 'int8'): (512, 8192, 1024),
+    (6, 2048, 8192, 3584, 'int8', 'int8'): (2048, 512, 3584),
+    (6, 2048, 896, 5120, 'int8', 'int8'): (512, 896, 5120),
+    (6, 256, 1024, 4096, 'int8', 'int8'): (256, 512, 4096),
+    (6, 256, 1024, 8192, 'int8', 'int8'): (256, 512, 8192),
+    (6, 256, 128, 8192, 'int8', 'int8'): (256, 128, 8192),
+    (6, 256, 1280, 8192, 'int8', 'int8'): (256, 1280, 2048),
+    (6, 256, 13824, 5120, 'int8', 'int8'): (256, 6912, 512),
+    (6, 256, 14336, 4096, 'int8', 'int8'): (256, 1024, 4096),
+    (6, 256, 1536, 1536, 'int8', 'int8'): (256, 768, 1536),
+    (6, 256, 1536, 8960, 'int8', 'int8'): (256, 256, 8960),
+    (6, 256, 1792, 5120, 'int8', 'int8'): (256, 896, 5120),
+    (6, 256, 17920, 1536, 'int8', 'int8'): (256, 3584, 1536),
+    (6, 256, 2048, 1536, 'int8', 'int8'): (256, 2048, 1536),
+    (6, 256, 28672, 4096, 'int8', 'int8'): (256, 1024, 4096),
+    (6, 256, 3584, 18944, 'int8', 'int8'): (256, 256, 18944),
+    (6, 256, 3584, 3584, 'int8', 'int8'): (256, 512, 3584),
+    (6, 256, 3584, 8192, 'int8', 'int8'): (256, 3584, 1024),
+    (6, 256, 37888, 3584, 'int8', 'int8'): (256, 9472, 512),
+    (6, 256, 4096, 14336, 'int8', 'int8'): (256, 2048, 2048),
+    (6, 256, 4096, 4096, 'int8', 'int8'): (256, 512, 4096),
+    (6, 256, 4608, 3584, 'int8', 'int8'): (256, 512, 3584),
+    (6, 256, 5120, 1280, 'int8', 'int8'): (256, 2560, 1280),
+    (6, 256, 5120, 3456, 'int8', 'int8'): (256, 640, 3456),
+    (6, 256, 5120, 640, 'int8', 'int8'): (256, 2560, 640),
+    (6, 256, 5120, 6912, 'int8', 'int8'): (256, 640, 6912),
+    (6, 256, 6144, 4096, 'int8', 'int8'): (256, 2048, 2048),
+    (6, 256, 6912, 5120, 'int8', 'int8'): (256, 768, 5120),
+    (6, 256, 7168, 8192, 'int8', 'int8'): (256, 512, 8192),
+    (6, 256, 8192, 1024, 'int8', 'int8'): (256, 2048, 1024),
+    (6, 256, 8192, 3584, 'int8', 'int8'): (256, 8192, 512),
+    (6, 256, 896, 5120, 'int8', 'int8'): (256, 896, 2560),
+    (6, 32, 1024, 4096, 'int8', 'int8'): (32, 512, 4096),
+    (6, 32, 1024, 8192, 'int8', 'int8'): (32, 256, 8192),
+    (6, 32, 128, 8192, 'int8', 'int8'): (32, 128, 4096),
+    (6, 32, 1280, 8192, 'int8', 'int8'): (32, 1280, 2048),
+    (6, 32, 13824, 5120, 'int8', 'int8'): (32, 512, 5120),
+    (6, 32, 14336, 4096, 'int8', 'int8'): (32, 896, 4096),
+    (6, 32, 1536, 1536, 'int8', 'int8'): (32, 768, 1536),
+    (6, 32, 1536, 8960, 'int8', 'int8'): (32, 1536, 1792),
+    (6, 32, 1792, 5120, 'int8', 'int8'): (32, 1792, 2560),
+    (6, 32, 17920, 1536, 'int8', 'int8'): (32, 1792, 1536),
+    (6, 32, 2048, 1536, 'int8', 'int8'): (32, 1024, 1536),
+    (6, 32, 28672, 4096, 'int8', 'int8'): (32, 896, 4096),
+    (6, 32, 3584, 18944, 'int8', 'int8'): (32, 256, 18944),
+    (6, 32, 3584, 3584, 'int8', 'int8'): (32, 896, 3584),
+    (6, 32, 3584, 8192, 'int8', 'int8'): (32, 512, 8192),
+    (6, 32, 37888, 3584, 'int8', 'int8'): (32, 37888, 128),
+    (6, 32, 4096, 14336, 'int8', 'int8'): (32, 256, 14336),
+    (6, 32, 4096, 4096, 'int8', 'int8'): (32, 1024, 4096),
+    (6, 32, 4608, 3584, 'int8', 'int8'): (32, 1536, 1792),
+    (6, 32, 5120, 1280, 'int8', 'int8'): (32, 2560, 1280),
+    (6, 32, 5120, 3456, 'int8', 'int8'): (32, 1024, 3456),
+    (6, 32, 5120, 640, 'int8', 'int8'): (32, 2560, 640),
+    (6, 32, 5120, 6912, 'int8', 'int8'): (32, 512, 6912),
+    (6, 32, 6144, 4096, 'int8', 'int8'): (32, 768, 4096),
+    (6, 32, 6912, 5120, 'int8', 'int8'): (32, 768, 5120),
+    (6, 32, 7168, 8192, 'int8', 'int8'): (32, 512, 8192),
+    (6, 32, 8192, 1024, 'int8', 'int8'): (32, 8192, 256),
+    (6, 32, 8192, 3584, 'int8', 'int8'): (32, 4096, 896),
+    (6, 32, 896, 5120, 'int8', 'int8'): (32, 896, 2560),
+    (6, 4096, 1024, 4096, 'int8', 'int8'): (2048, 512, 4096),
+    (6, 4096, 1024, 8192, 'int8', 'int8'): (4096, 1024, 1024),
+    (6, 4096, 128, 8192, 'int8', 'int8'): (2048, 128, 8192),
+    (6, 4096, 1280, 8192, 'int8', 'int8'): (1024, 1280, 8192),
+    (6, 4096, 13824, 5120, 'int8', 'int8'): (2048, 768, 5120),
+    (6, 4096, 14336, 4096, 'int8', 'int8'): (1024, 2048, 4096),
+    (6, 4096, 1536, 1536, 'int8', 'int8'): (512, 1536, 1536),
+    (6, 4096, 1536, 8960, 'int8', 'int8'): (1024, 1536, 4480),
+    (6, 4096, 1792, 5120, 'int8', 'int8'): (2048, 1792, 5120),
+    (6, 4096, 17920, 1536, 'int8', 'int8'): (256, 17920, 1536),
+    (6, 4096, 2048, 1536, 'int8', 'int8'): (512, 2048, 1536),
+    (6, 4096, 28672, 4096, 'int8', 'int8'): (2048, 1792, 4096),
+    (6, 4096, 3584, 18944, 'int8', 'int8'): (4096, 1792, 512),
+    (6, 4096, 3584, 3584, 'int8', 'int8'): (4096, 512, 3584),
+    (6, 4096, 3584, 8192, 'int8', 'int8'): (256, 3584, 8192),
+    (6, 4096, 37888, 3584, 'int8', 'int8'): (2048, 1024, 3584),
+    (6, 4096, 4096, 14336, 'int8', 'int8'): (1024, 4096, 1792),
+    (6, 4096, 4096, 4096, 'int8', 'int8'): (4096, 512, 4096),
+    (6, 4096, 4608, 3584, 'int8', 'int8'): (4096, 512, 3584),
+    (6, 4096, 5120, 1280, 'int8', 'int8'): (256, 5120, 1280),
+    (6, 4096, 5120, 3456, 'int8', 'int8'): (4096, 512, 3456),
+    (6, 4096, 5120, 640, 'int8', 'int8'): (512, 5120, 640),
+    (6, 4096, 5120, 6912, 'int8', 'int8'): (512, 5120, 6912),
+    (6, 4096, 6144, 4096, 'int8', 'int8'): (4096, 512, 4096),
+    (6, 4096, 6912, 5120, 'int8', 'int8'): (512, 6912, 5120),
+    (6, 4096, 7168, 8192, 'int8', 'int8'): (1024, 1792, 8192),
+    (6, 4096, 8192, 1024, 'int8', 'int8'): (4096, 512, 1024),
+    (6, 4096, 8192, 3584, 'int8', 'int8'): (4096, 512, 3584),
+    (6, 4096, 896, 5120, 'int8', 'int8'): (512, 896, 5120),
+    (6, 512, 1024, 4096, 'int8', 'int8'): (512, 512, 2048),
+    (6, 512, 1024, 8192, 'int8', 'int8'): (512, 512, 8192),
+    (6, 512, 128, 8192, 'int8', 'int8'): (256, 128, 8192),
+    (6, 512, 1280, 8192, 'int8', 'int8'): (512, 256, 8192),
+    (6, 512, 13824, 5120, 'int8', 'int8'): (512, 768, 5120),
+    (6, 512, 14336, 4096, 'int8', 'int8'): (512, 896, 4096),
+    (6, 512, 1536, 1536, 'int8', 'int8'): (512, 768, 1536),
+    (6, 512, 1536, 8960, 'int8', 'int8'): (512, 256, 8960),
+    (6, 512, 1792, 5120, 'int8', 'int8'): (512, 1792, 1280),
+    (6, 512, 17920, 1536, 'int8', 'int8'): (512, 1280, 1536),
+    (6, 512, 2048, 1536, 'int8', 'int8'): (256, 2048, 1536),
+    (6, 512, 28672, 4096, 'int8', 'int8'): (512, 1024, 4096),
+    (6, 512, 3584, 18944, 'int8', 'int8'): (512, 256, 18944),
+    (6, 512, 3584, 3584, 'int8', 'int8'): (512, 1792, 3584),
+    (6, 512, 3584, 8192, 'int8', 'int8'): (512, 3584, 2048),
+    (6, 512, 37888, 3584, 'int8', 'int8'): (512, 1024, 3584),
+    (6, 512, 4096, 14336, 'int8', 'int8'): (512, 512, 14336),
+    (6, 512, 4096, 4096, 'int8', 'int8'): (512, 1024, 4096),
+    (6, 512, 4608, 3584, 'int8', 'int8'): (512, 768, 3584),
+    (6, 512, 5120, 1280, 'int8', 'int8'): (256, 5120, 1280),
+    (6, 512, 5120, 3456, 'int8', 'int8'): (512, 1280, 3456),
+    (6, 512, 5120, 640, 'int8', 'int8'): (512, 2560, 640),
+    (6, 512, 5120, 6912, 'int8', 'int8'): (512, 512, 6912),
+    (6, 512, 6144, 4096, 'int8', 'int8'): (512, 768, 4096),
+    (6, 512, 6912, 5120, 'int8', 'int8'): (512, 768, 5120),
+    (6, 512, 7168, 8192, 'int8', 'int8'): (512, 1792, 4096),
+    (6, 512, 8192, 1024, 'int8', 'int8'): (512, 2048, 1024),
+    (6, 512, 8192, 3584, 'int8', 'int8'): (512, 1024, 3584),
+    (6, 512, 896, 5120, 'int8', 'int8'): (512, 896, 2560),
+    (6, 64, 1024, 4096, 'int8', 'int8'): (64, 1024, 1024),
+    (6, 64, 1024, 8192, 'int8', 'int8'): (64, 1024, 2048),
+    (6, 64, 128, 8192, 'int8', 'int8'): (64, 128, 4096),
+    (6, 64, 1280, 8192, 'int8', 'int8'): (64, 1280, 2048),
+    (6, 64, 13824, 5120, 'int8', 'int8'): (64, 13824, 256),
+    (6, 64, 14336, 4096, 'int8', 'int8'): (64, 7168, 512),
+    (6, 64, 1536, 1536, 'int8', 'int8'): (64, 256, 1536),
+    (6, 64, 1536, 8960, 'int8', 'int8'): (64, 256, 8960),
+    (6, 64, 1792, 5120, 'int8', 'int8'): (64, 896, 5120),
+    (6, 64, 17920, 1536, 'int8', 'int8'): (64, 1792, 1536),
+    (6, 64, 2048, 1536, 'int8', 'int8'): (64, 2048, 768),
+    (6, 64, 28672, 4096, 'int8', 'int8'): (64, 1024, 4096),
+    (6, 64, 3584, 18944, 'int8', 'int8'): (64, 256, 18944),
+    (6, 64, 3584, 3584, 'int8', 'int8'): (64, 896, 3584),
+    (6, 64, 3584, 8192, 'int8', 'int8'): (64, 512, 8192),
+    (6, 64, 37888, 3584, 'int8', 'int8'): (64, 18944, 256),
+    (6, 64, 4096, 14336, 'int8', 'int8'): (64, 256, 14336),
+    (6, 64, 4096, 4096, 'int8', 'int8'): (64, 4096, 1024),
+    (6, 64, 4608, 3584, 'int8', 'int8'): (64, 768, 3584),
+    (6, 64, 5120, 1280, 'int8', 'int8'): (64, 2560, 1280),
+    (6, 64, 5120, 3456, 'int8', 'int8'): (64, 1024, 3456),
+    (6, 64, 5120, 640, 'int8', 'int8'): (64, 2560, 640),
+    (6, 64, 5120, 6912, 'int8', 'int8'): (64, 512, 6912),
+    (6, 64, 6144, 4096, 'int8', 'int8'): (64, 1536, 4096),
+    (6, 64, 6912, 5120, 'int8', 'int8'): (64, 768, 5120),
+    (6, 64, 7168, 8192, 'int8', 'int8'): (64, 1024, 8192),
+    (6, 64, 8192, 1024, 'int8', 'int8'): (64, 2048, 1024),
+    (6, 64, 8192, 3584, 'int8', 'int8'): (64, 1024, 3584),
+    (6, 64, 896, 5120, 'int8', 'int8'): (64, 896, 2560),
+    (6, 8192, 1024, 4096, 'int8', 'int8'): (2048, 1024, 4096),
+    (6, 8192, 1024, 8192, 'int8', 'int8'): (512, 1024, 8192),
+    (6, 8192, 128, 8192, 'int8', 'int8'): (256, 128, 8192),
+    (6, 8192, 1280, 8192, 'int8', 'int8'): (512, 1280, 8192),
+    (6, 8192, 13824, 5120, 'int8', 'int8'): (2048, 1536, 5120),
+    (6, 8192, 14336, 4096, 'int8', 'int8'): (2048, 1792, 4096),
+    (6, 8192, 1536, 1536, 'int8', 'int8'): (1024, 1536, 1536),
+    (6, 8192, 1536, 8960, 'int8', 'int8'): (256, 1536, 8960),
+    (6, 8192, 1792, 5120, 'int8', 'int8'): (1024, 1792, 2560),
+    (6, 8192, 17920, 1536, 'int8', 'int8'): (512, 17920, 1536),
+    (6, 8192, 2048, 1536, 'int8', 'int8'): (512, 2048, 1536),
+    (6, 8192, 28672, 4096, 'int8', 'int8'): (2048, 1792, 4096),
+    (6, 8192, 3584, 18944, 'int8', 'int8'): (4096, 1792, 512),
+    (6, 8192, 3584, 3584, 'int8', 'int8'): (2048, 1792, 3584),
+    (6, 8192, 3584, 8192, 'int8', 'int8'): (256, 3584, 8192),
+    (6, 8192, 37888, 3584, 'int8', 'int8'): (4096, 512, 3584),
+    (6, 8192, 4096, 14336, 'int8', 'int8'): (1024, 2048, 3584),
+    (6, 8192, 4096, 4096, 'int8', 'int8'): (1024, 2048, 4096),
+    (6, 8192, 4608, 3584, 'int8', 'int8'): (2048, 1536, 3584),
+    (6, 8192, 5120, 1280, 'int8', 'int8'): (2048, 1280, 1280),
+    (6, 8192, 5120, 3456, 'int8', 'int8'): (1024, 5120, 3456),
+    (6, 8192, 5120, 640, 'int8', 'int8'): (512, 5120, 640),
+    (6, 8192, 5120, 6912, 'int8', 'int8'): (512, 5120, 6912),
+    (6, 8192, 6144, 4096, 'int8', 'int8'): (1024, 6144, 4096),
+    (6, 8192, 6912, 5120, 'int8', 'int8'): (512, 6912, 5120),
+    (6, 8192, 7168, 8192, 'int8', 'int8'): (256, 1792, 8192),
+    (6, 8192, 8192, 1024, 'int8', 'int8'): (512, 8192, 1024),
+    (6, 8192, 8192, 3584, 'int8', 'int8'): (1024, 4096, 3584),
+    (6, 8192, 896, 5120, 'int8', 'int8'): (2048, 896, 5120),
+    (7, 1024, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'):
+    (1024, 1792, 4096),
+    (7, 1024, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'):
+    (1024, 1280, 5120),
+    (7, 1024, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'):
+    (512, 1024, 8192),
+    (7, 1024, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'):
+    (1024, 3072, 1024),
+    (7, 1024, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1024,
+                                                              2048),
+    (7, 1024, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                              1024),
+    (7, 1024, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 1024, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 1024, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              4096),
+    (7, 1024, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              5120),
+    (7, 1024, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 5120, 2048),
+    (7, 1024, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                               2048),
+    (7, 1024, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1024,
+                                                              4096),
+    (7, 128, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 3584, 4096),
+    (7, 128, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 3200, 5120),
+    (7, 128, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1792, 8192),
+    (7, 128, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1536, 4096),
+    (7, 128, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1024, 2048),
+    (7, 128, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1024, 7168),
+    (7, 128, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1024,
+                                                              12800),
+    (7, 128, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 512, 25600),
+    (7, 128, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 2560, 4096),
+    (7, 128, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1280, 5120),
+    (7, 128, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 5120, 2048),
+    (7, 128, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 1024,
+                                                              14336),
+    (7, 128, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (128, 4096, 4096),
+    (7, 16, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 3584, 4096),
+    (7, 16, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 3200, 5120),
+    (7, 16, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 2048, 8192),
+    (7, 16, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 1536, 4096),
+    (7, 16, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 4096, 1024),
+    (7, 16, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 4096, 1792),
+    (7, 16, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 1024, 12800),
+    (7, 16, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 640, 25600),
+    (7, 16, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 1280, 4096),
+    (7, 16, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 1280, 5120),
+    (7, 16, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 1280, 8192),
+    (7, 16, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 8192, 1024),
+    (7, 16, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (16, 2048, 4096),
+    (7, 2048, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                               4096),
+    (7, 2048, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2560,
+                                                               5120),
+    (7, 2048, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                               1024),
+    (7, 2048, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1536,
+                                                              4096),
+    (7, 2048, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                              2048),
+    (7, 2048, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                              1024),
+    (7, 2048, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 2048, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 2048, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              4096),
+    (7, 2048, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                              2560),
+    (7, 2048, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                              2048),
+    (7, 2048, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                               1024),
+    (7, 2048, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                              4096),
+    (7, 256, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 2048, 4096),
+    (7, 256, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 2560, 5120),
+    (7, 256, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 7168, 2048),
+    (7, 256, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 3072, 2048),
+    (7, 256, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 2048, 2048),
+    (7, 256, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 1024, 7168),
+    (7, 256, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 5120, 2560),
+    (7, 256, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 5120, 2560),
+    (7, 256, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 1280, 4096),
+    (7, 256, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 1280, 5120),
+    (7, 256, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 1280, 8192),
+    (7, 256, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 8192, 1792),
+    (7, 256, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (256, 2048, 4096),
+    (7, 32, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 1792, 4096),
+    (7, 32, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 2560, 5120),
+    (7, 32, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 2048, 8192),
+    (7, 32, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 3072, 2048),
+    (7, 32, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 1024, 2048),
+    (7, 32, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 2048, 7168),
+    (7, 32, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 1280, 12800),
+    (7, 32, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 640, 25600),
+    (7, 32, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 1280, 4096),
+    (7, 32, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 1280, 5120),
+    (7, 32, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 2560, 4096),
+    (7, 32, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 2048, 7168),
+    (7, 32, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (32, 4096, 4096),
+    (7, 4096, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 3584,
+                                                               4096),
+    (7, 4096, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                               5120),
+    (7, 4096, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1024,
+                                                               8192),
+    (7, 4096, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 3072, 4096),
+    (7, 4096, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 4096, 2048),
+    (7, 4096, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                              1024),
+    (7, 4096, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 4096, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                               2560),
+    (7, 4096, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              4096),
+    (7, 4096, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              5120),
+    (7, 4096, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2560,
+                                                              2048),
+    (7, 4096, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                               1024),
+    (7, 4096, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1024,
+                                                              4096),
+    (7, 512, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2048, 4096),
+    (7, 512, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2560, 5120),
+    (7, 512, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 512, 8192),
+    (7, 512, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 3072, 1024),
+    (7, 512, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2048, 2048),
+    (7, 512, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1024, 7168),
+    (7, 512, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 512, 12800),
+    (7, 512, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2560, 2560),
+    (7, 512, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2560, 2048),
+    (7, 512, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1280, 5120),
+    (7, 512, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1024, 8192),
+    (7, 512, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 4096, 2048),
+    (7, 512, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1024, 4096),
+    (7, 64, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 14336, 1024),
+    (7, 64, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 6400, 2560),
+    (7, 64, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1792, 8192),
+    (7, 64, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1536, 4096),
+    (7, 64, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 2048, 2048),
+    (7, 64, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1024, 7168),
+    (7, 64, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1024, 12800),
+    (7, 64, 5120, 25600, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 640, 25600),
+    (7, 64, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 2560, 4096),
+    (7, 64, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1280, 5120),
+    (7, 64, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1280, 8192),
+    (7, 64, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 1024, 14336),
+    (7, 64, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (64, 2048, 4096),
+    (7, 8192, 14336, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1024,
+                                                               4096),
+    (7, 8192, 25600, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 1280,
+                                                               5120),
+    (7, 8192, 28672, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                               2048),
+    (7, 8192, 3072, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1536,
+                                                              4096),
+    (7, 8192, 4096, 2048, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 4096, 2048),
+    (7, 8192, 4096, 7168, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                              1024),
+    (7, 8192, 5120, 12800, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 2560,
+                                                               3200),
+    (7, 8192, 5120, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              4096),
+    (7, 8192, 5120, 5120, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 1280,
+                                                              5120),
+    (7, 8192, 5120, 8192, 'float8_e4m3fn', 'float8_e4m3fn'): (512, 5120, 2048),
+    (7, 8192, 8192, 14336, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 4096,
+                                                               1024),
+    (7, 8192, 8192, 4096, 'float8_e4m3fn', 'float8_e4m3fn'): (1024, 2048,
+                                                              4096),
+    # go/keep-sorted end
+}
+
+TUNED_BLOCK_SIZES: dict[TunedKey, TunedValue] = {
+    TunedKey(*key): TunedValue(*value)
+    for key, value in TUNED_BLOCK_SIZES_RAW.items()
+}
+
+DEVICE_VMEM_LIMIT = {6: 96 * 1024 * 1024, 7: 48 * 1024 * 1024}
+
+
+def get_device_vmem_limit() -> int:
+    tpu_version = get_tpu_version()
+    if tpu_version not in DEVICE_VMEM_LIMIT:
+        logger.warning_once(
+            'VMEM limit for TPU version %d not found. Using default VMEM limit '
+            'of 96MiB', tpu_version)
+        return 96 * 1024 * 1024
+    return DEVICE_VMEM_LIMIT[tpu_version]
+
+
+def get_tpu_version() -> int:
+    """Returns the numeric version of the TPU, or -1 if not on TPU."""
+    kind = jax.devices()[0].device_kind
+    match = re.match(r'^TPU[^\d]*(\d+)', kind)
+    if match is None:
+        return -1
+    return int(match.group(1))
+
+
+def get_key(
+    n_batch: int,
+    n_out: int,
+    n_in: int,
+    x_q_dtype: str,
+    w_q_dtype: str,
+) -> TunedKey:
+    """Returns the key for the given parameters."""
+    return TunedKey(
+        get_tpu_version(),
+        n_batch,
+        n_out,
+        n_in,
+        x_q_dtype,
+        w_q_dtype,
+    )
+
+
+def get_tuned_block_sizes(
+    n_batch: int,
+    n_out: int,
+    n_in: int,
+    x_q_dtype: str,
+    w_q_dtype: str,
+) -> TunedValue:
+    """Retrieve the tuned block sizes for the given parameters.
+
+  Args:
+      n_batch: The batch size.
+      n_out: The number of output features.
+      n_in: The number of input features.
+      x_q_dtype: The data type of the activation ('int8' or 'float8_e4m3fn').
+      w_q_dtype: The data type of the weight ('int8' or 'float8_e4m3fn').
+
+  Returns:
+      tuple: A tuple containing the batch_block_size, out_block_size, and
+      in_block_size.
+  """
+    key = get_key(
+        n_batch,
+        n_out,
+        n_in,
+        x_q_dtype,
+        w_q_dtype,
+    )
+    tuned_value = TUNED_BLOCK_SIZES.get(key)
+    if tuned_value is None:
+        logger.warning_once(
+            'Couldn`t find tuned sizes for the quantized matmul kernel with %s',
+            key)
+        return TunedValue(128, 128, 128)
+    else:
+        return tuned_value