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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (260) hide show
  1. tests/__init__.py +13 -0
  2. tests/core/__init__.py +13 -0
  3. tests/core/test_core_tpu.py +513 -0
  4. tests/core/test_disagg_executor.py +60 -0
  5. tests/core/test_disagg_utils.py +67 -0
  6. tests/core/test_dp_scheduler.py +724 -0
  7. tests/core/test_init.py +63 -0
  8. tests/distributed/__init__.py +13 -0
  9. tests/distributed/test_distributed_utils.py +120 -0
  10. tests/distributed/test_tpu_connector.py +478 -0
  11. tests/e2e/__init__.py +13 -0
  12. tests/e2e/test_async_scheduler.py +211 -0
  13. tests/e2e/test_data_parallel.py +393 -0
  14. tests/e2e/test_local_disagg.py +257 -0
  15. tests/e2e/test_model_loader.py +268 -0
  16. tests/e2e/test_multi_modal_inference.py +111 -0
  17. tests/e2e/test_pipeline_parallel.py +265 -0
  18. tests/e2e/test_runai_model_streamer_loader.py +104 -0
  19. tests/e2e/test_sampling_params.py +269 -0
  20. tests/e2e/test_speculative_decoding.py +291 -0
  21. tests/e2e/test_structured_decoding.py +46 -0
  22. tests/executors/__init__.py +13 -0
  23. tests/executors/test_ray_distributed_executor.py +199 -0
  24. tests/experimental/__init__.py +13 -0
  25. tests/experimental/test_llama3_jax_stashed.py +208 -0
  26. tests/kernels/__init__.py +13 -0
  27. tests/kernels/collectives/__init__.py +13 -0
  28. tests/kernels/collectives/all_gather_matmul_kernel_test.py +69 -0
  29. tests/kernels/fused_moe_v1_test.py +388 -0
  30. tests/kernels/gmm_test.py +205 -0
  31. tests/kernels/mla_v1_test.py +498 -0
  32. tests/kernels/quantized_matmul_kernel_test.py +159 -0
  33. tests/kernels/ragged_kv_cache_update_v2_test.py +248 -0
  34. tests/kernels/ragged_paged_attention_kernel_v2_test.py +414 -0
  35. tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py +565 -0
  36. tests/kernels/ragged_paged_attention_kernel_v3_test.py +520 -0
  37. tests/layers/__init__.py +13 -0
  38. tests/layers/common/__init__.py +13 -0
  39. tests/layers/common/test_attention_interface.py +156 -0
  40. tests/layers/common/test_quantization.py +149 -0
  41. tests/layers/jax/__init__.py +13 -0
  42. tests/layers/jax/attention/__init__.py +13 -0
  43. tests/layers/jax/attention/test_common_attention.py +103 -0
  44. tests/layers/jax/attention/test_deepseek_v3_attention.py +233 -0
  45. tests/layers/jax/attention/test_llama4_attention.py +135 -0
  46. tests/layers/jax/moe/__init__.py +13 -0
  47. tests/layers/jax/moe/test_deepseek_moe.py +235 -0
  48. tests/layers/jax/sample/__init__.py +13 -0
  49. tests/layers/jax/sample/test_rejection_sampler.py +1624 -0
  50. tests/layers/jax/sample/test_sampling.py +115 -0
  51. tests/layers/jax/sample/test_sampling_metadata.py +254 -0
  52. tests/layers/jax/test_layers.py +155 -0
  53. tests/layers/jax/test_qwix.py +969 -0
  54. tests/layers/jax/test_rope.py +93 -0
  55. tests/layers/jax/test_sharding.py +159 -0
  56. tests/layers/jax/test_transformer_block.py +152 -0
  57. tests/layers/vllm/__init__.py +13 -0
  58. tests/layers/vllm/test_attention.py +363 -0
  59. tests/layers/vllm/test_awq.py +405 -0
  60. tests/layers/vllm/test_compressed_tensors_moe.py +202 -0
  61. tests/layers/vllm/test_compressed_tensors_w8a8_fp8.py +403 -0
  62. tests/layers/vllm/test_compressed_tensors_w8a8_int8.py +426 -0
  63. tests/layers/vllm/test_fp8.py +17 -0
  64. tests/layers/vllm/test_mxfp4.py +297 -0
  65. tests/layers/vllm/test_unquantized.py +621 -0
  66. tests/layers/vllm/utils.py +72 -0
  67. tests/lora/__init__.py +13 -0
  68. tests/lora/conftest.py +46 -0
  69. tests/lora/test_bgmv.py +57 -0
  70. tests/lora/test_layers.py +666 -0
  71. tests/lora/test_lora.py +147 -0
  72. tests/lora/test_lora_perf.py +67 -0
  73. tests/lora/utils.py +88 -0
  74. tests/models/__init__.py +13 -0
  75. tests/models/common/__init__.py +13 -0
  76. tests/models/common/test_model_loader.py +455 -0
  77. tests/models/jax/__init__.py +13 -0
  78. tests/models/jax/test_deepseek_v3.py +401 -0
  79. tests/models/jax/test_llama3.py +184 -0
  80. tests/models/jax/test_llama4.py +298 -0
  81. tests/models/jax/test_llama_eagle3.py +197 -0
  82. tests/models/jax/test_llama_guard_4.py +242 -0
  83. tests/models/jax/test_qwen2.py +172 -0
  84. tests/models/jax/test_qwen2_5_vl.py +606 -0
  85. tests/models/jax/test_qwen3.py +169 -0
  86. tests/models/jax/test_weight_loading.py +180 -0
  87. tests/models/jax/utils/__init__.py +13 -0
  88. tests/models/jax/utils/test_multi_modal_utils.py +212 -0
  89. tests/platforms/__init__.py +13 -0
  90. tests/platforms/test_tpu_platform.py +54 -0
  91. tests/runner/__init__.py +13 -0
  92. tests/runner/test_block_table.py +395 -0
  93. tests/runner/test_input_batch.py +226 -0
  94. tests/runner/test_kv_cache.py +220 -0
  95. tests/runner/test_kv_cache_manager.py +498 -0
  96. tests/runner/test_multimodal_manager.py +429 -0
  97. tests/runner/test_persistent_batch_manager.py +84 -0
  98. tests/runner/test_speculative_decoding_manager.py +368 -0
  99. tests/runner/test_structured_decoding_manager.py +220 -0
  100. tests/runner/test_tpu_runner.py +202 -0
  101. tests/runner/test_tpu_runner_dp.py +1033 -0
  102. tests/runner/test_tpu_runner_mesh.py +200 -0
  103. tests/runner/test_utils.py +411 -0
  104. tests/spec_decode/__init__.py +13 -0
  105. tests/spec_decode/test_eagle3.py +311 -0
  106. tests/test_base.py +215 -0
  107. tests/test_envs.py +280 -0
  108. tests/test_tpu_info.py +134 -0
  109. tests/test_utils.py +193 -0
  110. tests/worker/__init__.py +13 -0
  111. tests/worker/tpu_worker_test.py +414 -0
  112. tpu_inference/__init__.py +67 -0
  113. tpu_inference/core/__init__.py +13 -0
  114. tpu_inference/core/core_tpu.py +786 -0
  115. tpu_inference/core/disagg_executor.py +118 -0
  116. tpu_inference/core/disagg_utils.py +49 -0
  117. tpu_inference/core/sched/__init__.py +13 -0
  118. tpu_inference/core/sched/dp_scheduler.py +814 -0
  119. tpu_inference/distributed/__init__.py +13 -0
  120. tpu_inference/distributed/jax_parallel_state.py +81 -0
  121. tpu_inference/distributed/tpu_connector.py +732 -0
  122. tpu_inference/distributed/utils.py +112 -0
  123. tpu_inference/env_override.py +9 -0
  124. tpu_inference/envs.py +191 -0
  125. tpu_inference/executors/__init__.py +13 -0
  126. tpu_inference/executors/ray_distributed_executor.py +399 -0
  127. tpu_inference/experimental/__init__.py +13 -0
  128. tpu_inference/experimental/llama3_jax_stashed.py +272 -0
  129. tpu_inference/kernels/__init__.py +13 -0
  130. tpu_inference/kernels/collectives/__init__.py +13 -0
  131. tpu_inference/kernels/collectives/all_gather_matmul.py +741 -0
  132. tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py +65 -0
  133. tpu_inference/kernels/collectives/util.py +47 -0
  134. tpu_inference/kernels/flash_attention/__init__.py +13 -0
  135. tpu_inference/kernels/flash_attention/kernel.py +772 -0
  136. tpu_inference/kernels/fused_moe/__init__.py +13 -0
  137. tpu_inference/kernels/fused_moe/v1/__init__.py +13 -0
  138. tpu_inference/kernels/fused_moe/v1/kernel.py +1612 -0
  139. tpu_inference/kernels/megablox/__init__.py +13 -0
  140. tpu_inference/kernels/megablox/common.py +54 -0
  141. tpu_inference/kernels/megablox/gmm.py +646 -0
  142. tpu_inference/kernels/mla/__init__.py +13 -0
  143. tpu_inference/kernels/mla/v1/__init__.py +13 -0
  144. tpu_inference/kernels/mla/v1/kernel.py +1340 -0
  145. tpu_inference/kernels/quantized_matmul/__init__.py +13 -0
  146. tpu_inference/kernels/quantized_matmul/kernel.py +456 -0
  147. tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py +609 -0
  148. tpu_inference/kernels/quantized_matmul/util.py +58 -0
  149. tpu_inference/kernels/ragged_paged_attention/__init__.py +13 -0
  150. tpu_inference/kernels/ragged_paged_attention/v2/__init__.py +13 -0
  151. tpu_inference/kernels/ragged_paged_attention/v2/kernel.py +876 -0
  152. tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py +288 -0
  153. tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py +1482 -0
  154. tpu_inference/kernels/ragged_paged_attention/v3/__init__.py +13 -0
  155. tpu_inference/kernels/ragged_paged_attention/v3/kernel.py +1594 -0
  156. tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py +1586 -0
  157. tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py +4460 -0
  158. tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py +548 -0
  159. tpu_inference/kernels/ragged_paged_attention/v3/util.py +65 -0
  160. tpu_inference/layers/__init__.py +13 -0
  161. tpu_inference/layers/common/__init__.py +13 -0
  162. tpu_inference/layers/common/attention_interface.py +403 -0
  163. tpu_inference/layers/common/attention_metadata.py +48 -0
  164. tpu_inference/layers/common/binary_search.py +295 -0
  165. tpu_inference/layers/common/quant_methods.py +23 -0
  166. tpu_inference/layers/common/quantization.py +270 -0
  167. tpu_inference/layers/common/sharding.py +600 -0
  168. tpu_inference/layers/jax/__init__.py +13 -0
  169. tpu_inference/layers/jax/attention/__init__.py +13 -0
  170. tpu_inference/layers/jax/attention/attention.py +268 -0
  171. tpu_inference/layers/jax/attention/deepseek_v3_attention.py +547 -0
  172. tpu_inference/layers/jax/attention/gpt_oss_attention.py +275 -0
  173. tpu_inference/layers/jax/attention/llama4_attention.py +167 -0
  174. tpu_inference/layers/jax/base.py +165 -0
  175. tpu_inference/layers/jax/constants.py +101 -0
  176. tpu_inference/layers/jax/layers.py +315 -0
  177. tpu_inference/layers/jax/misc.py +30 -0
  178. tpu_inference/layers/jax/moe/__init__.py +13 -0
  179. tpu_inference/layers/jax/moe/deepseek_v3_moe.py +615 -0
  180. tpu_inference/layers/jax/moe/gpt_oss_moe.py +199 -0
  181. tpu_inference/layers/jax/moe/moe.py +249 -0
  182. tpu_inference/layers/jax/pp_utils.py +53 -0
  183. tpu_inference/layers/jax/rope.py +294 -0
  184. tpu_inference/layers/jax/rope_interface.py +228 -0
  185. tpu_inference/layers/jax/sample/__init__.py +13 -0
  186. tpu_inference/layers/jax/sample/rejection_sampler.py +528 -0
  187. tpu_inference/layers/jax/sample/sampling.py +110 -0
  188. tpu_inference/layers/jax/sample/sampling_metadata.py +90 -0
  189. tpu_inference/layers/jax/transformer_block.py +121 -0
  190. tpu_inference/layers/vllm/__init__.py +13 -0
  191. tpu_inference/layers/vllm/attention.py +221 -0
  192. tpu_inference/layers/vllm/fused_moe.py +502 -0
  193. tpu_inference/layers/vllm/linear_common.py +221 -0
  194. tpu_inference/layers/vllm/quantization/__init__.py +55 -0
  195. tpu_inference/layers/vllm/quantization/awq.py +221 -0
  196. tpu_inference/layers/vllm/quantization/common.py +124 -0
  197. tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py +13 -0
  198. tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py +135 -0
  199. tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py +266 -0
  200. tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py +13 -0
  201. tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +222 -0
  202. tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +150 -0
  203. tpu_inference/layers/vllm/quantization/fp8.py +118 -0
  204. tpu_inference/layers/vllm/quantization/mxfp4.py +396 -0
  205. tpu_inference/layers/vllm/quantization/unquantized.py +416 -0
  206. tpu_inference/layers/vllm/sharding.py +244 -0
  207. tpu_inference/logger.py +10 -0
  208. tpu_inference/lora/__init__.py +13 -0
  209. tpu_inference/lora/torch_lora_ops.py +98 -0
  210. tpu_inference/lora/torch_punica_tpu.py +310 -0
  211. tpu_inference/models/__init__.py +13 -0
  212. tpu_inference/models/common/__init__.py +13 -0
  213. tpu_inference/models/common/model_loader.py +520 -0
  214. tpu_inference/models/jax/__init__.py +13 -0
  215. tpu_inference/models/jax/deepseek_v3.py +978 -0
  216. tpu_inference/models/jax/gpt_oss.py +508 -0
  217. tpu_inference/models/jax/jax_intermediate_tensor.py +93 -0
  218. tpu_inference/models/jax/llama3.py +436 -0
  219. tpu_inference/models/jax/llama4.py +643 -0
  220. tpu_inference/models/jax/llama_eagle3.py +350 -0
  221. tpu_inference/models/jax/llama_guard_4.py +375 -0
  222. tpu_inference/models/jax/qwen2.py +390 -0
  223. tpu_inference/models/jax/qwen2_5_vl.py +1232 -0
  224. tpu_inference/models/jax/qwen3.py +318 -0
  225. tpu_inference/models/jax/utils/__init__.py +13 -0
  226. tpu_inference/models/jax/utils/file_utils.py +110 -0
  227. tpu_inference/models/jax/utils/multi_modal_utils.py +177 -0
  228. tpu_inference/models/jax/utils/qwix/__init__.py +13 -0
  229. tpu_inference/models/jax/utils/qwix/qwix_utils.py +713 -0
  230. tpu_inference/models/jax/utils/weight_utils.py +621 -0
  231. tpu_inference/models/vllm/__init__.py +13 -0
  232. tpu_inference/models/vllm/vllm_model_wrapper.py +307 -0
  233. tpu_inference/models/vllm/vllm_model_wrapper_context.py +59 -0
  234. tpu_inference/platforms/__init__.py +16 -0
  235. tpu_inference/platforms/tpu_platform.py +258 -0
  236. tpu_inference/runner/__init__.py +13 -0
  237. tpu_inference/runner/block_table.py +122 -0
  238. tpu_inference/runner/compilation_manager.py +890 -0
  239. tpu_inference/runner/input_batch.py +435 -0
  240. tpu_inference/runner/kv_cache.py +166 -0
  241. tpu_inference/runner/kv_cache_manager.py +508 -0
  242. tpu_inference/runner/lora_utils.py +106 -0
  243. tpu_inference/runner/multimodal_manager.py +231 -0
  244. tpu_inference/runner/persistent_batch_manager.py +296 -0
  245. tpu_inference/runner/speculative_decoding_manager.py +262 -0
  246. tpu_inference/runner/structured_decoding_manager.py +101 -0
  247. tpu_inference/runner/tpu_runner.py +1768 -0
  248. tpu_inference/runner/utils.py +426 -0
  249. tpu_inference/spec_decode/__init__.py +13 -0
  250. tpu_inference/spec_decode/jax/__init__.py +13 -0
  251. tpu_inference/spec_decode/jax/eagle3.py +430 -0
  252. tpu_inference/tpu_info.py +92 -0
  253. tpu_inference/utils.py +345 -0
  254. tpu_inference/worker/__init__.py +13 -0
  255. tpu_inference/worker/tpu_worker.py +468 -0
  256. tpu_inference-0.12.0.dev20251222.dist-info/METADATA +106 -0
  257. tpu_inference-0.12.0.dev20251222.dist-info/RECORD +260 -0
  258. tpu_inference-0.12.0.dev20251222.dist-info/WHEEL +5 -0
  259. tpu_inference-0.12.0.dev20251222.dist-info/licenses/LICENSE +201 -0
  260. tpu_inference-0.12.0.dev20251222.dist-info/top_level.txt +2 -0
tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py ADDED
@@ -0,0 +1,548 @@
1
+ # Copyright 2025 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """Auto-tuned block sizes for ragged paged attention."""
15
+
16
+ import jax.numpy as jnp
17
+
18
+ from tpu_inference.kernels.ragged_paged_attention.v3.util import (
19
+ align_to, get_dtype_packing, get_tpu_version, next_power_of_2)
20
+ from tpu_inference.logger import init_logger
21
+ from tpu_inference.utils import get_device_name
22
+
23
+ logger = init_logger(__name__)
24
+
25
+ # key
26
+ # - device_name
27
+ # - page_size
28
+ # - q_{q_dtype_name}_kv_{kv_dtype_name}
29
+ # - q_head-{num_q_heads}_kv_head-{num_kv_heads}-_head-{head_dim}
30
+ # - max_model_len-{max_model_len}-sw-{sliding_window}
31
+ # value:
32
+ # - (num_kv_pages_per_block, num_queries_per_block)
33
+ TUNED_BLOCK_SIZES = {
34
+ 'TPU v7': {
35
+ 256: {
36
+ 'q_bfloat16_kv_bfloat16': {
37
+ 'q_head-16_kv_head-2_head-64': {
38
+ 'max_model_len-2048-sw-None': (8, 32),
39
+ 'max_model_len-2048-sw-128': (1, 32),
40
+ 'max_model_len-4096-sw-None': (16, 32),
41
+ 'max_model_len-4096-sw-128': (1, 32),
42
+ 'max_model_len-8192-sw-None': (16, 32),
43
+ 'max_model_len-8192-sw-128': (1, 32),
44
+ 'max_model_len-256-sw-None': (1, 32),
45
+ 'max_model_len-256-sw-128': (1, 32),
46
+ 'max_model_len-512-sw-None': (2, 32),
47
+ 'max_model_len-512-sw-128': (1, 32),
48
+ 'max_model_len-1024-sw-None': (4, 32),
49
+ 'max_model_len-1024-sw-128': (1, 64),
50
+ },
51
+ 'q_head-64_kv_head-8_head-64': {
52
+ 'max_model_len-8192-sw-128': (1, 16),
53
+ 'max_model_len-512-sw-None': (2, 16),
54
+ 'max_model_len-256-sw-None': (1, 16),
55
+ 'max_model_len-256-sw-128': (1, 16),
56
+ 'max_model_len-512-sw-128': (1, 16),
57
+ 'max_model_len-1024-sw-None': (4, 16),
58
+ 'max_model_len-1024-sw-128': (1, 16),
59
+ 'max_model_len-2048-sw-None': (8, 32),
60
+ 'max_model_len-2048-sw-128': (1, 16),
61
+ 'max_model_len-4096-sw-None': (16, 16),
62
+ 'max_model_len-4096-sw-128': (1, 16),
63
+ 'max_model_len-8192-sw-None': (16, 16),
64
+ },
65
+ 'q_head-32_kv_head-4_head-64': {
66
+ 'max_model_len-256-sw-None': (1, 32),
67
+ 'max_model_len-256-sw-128': (1, 32),
68
+ 'max_model_len-512-sw-None': (2, 32),
69
+ 'max_model_len-512-sw-128': (1, 32),
70
+ 'max_model_len-4096-sw-None': (16, 32),
71
+ 'max_model_len-1024-sw-None': (4, 32),
72
+ 'max_model_len-1024-sw-128': (1, 32),
73
+ 'max_model_len-2048-sw-None': (8, 32),
74
+ 'max_model_len-2048-sw-128': (1, 32),
75
+ 'max_model_len-4096-sw-128': (1, 32),
76
+ 'max_model_len-8192-sw-None': (16, 32),
77
+ 'max_model_len-8192-sw-128': (1, 32),
78
+ },
79
+ 'q_head-8_kv_head-2_head-64': {
80
+ 'max_model_len-256-sw-None': (1, 64),
81
+ 'max_model_len-256-sw-128': (1, 64),
82
+ 'max_model_len-512-sw-None': (2, 64),
83
+ 'max_model_len-512-sw-128': (1, 64),
84
+ 'max_model_len-1024-sw-None': (4, 64),
85
+ 'max_model_len-1024-sw-128': (1, 64),
86
+ 'max_model_len-2048-sw-None': (8, 32),
87
+ 'max_model_len-2048-sw-128': (1, 128),
88
+ 'max_model_len-4096-sw-None': (16, 32),
89
+ 'max_model_len-4096-sw-128': (1, 128),
90
+ 'max_model_len-8192-sw-None': (16, 32),
91
+ 'max_model_len-8192-sw-128': (1, 64),
92
+ },
93
+ }
94
+ },
95
+ 128: {
96
+ 'q_bfloat16_kv_bfloat16': {
97
+ 'q_head-64_kv_head-8_head-64': {
98
+ 'max_model_len-4096-sw-128': (2, 16),
99
+ 'max_model_len-8192-sw-None': (32, 16),
100
+ 'max_model_len-1024-sw-None': (8, 16),
101
+ 'max_model_len-8192-sw-128': (2, 16),
102
+ 'max_model_len-512-sw-None': (4, 16),
103
+ 'max_model_len-2048-sw-128': (2, 16),
104
+ 'max_model_len-256-sw-None': (2, 16),
105
+ 'max_model_len-128-sw-None': (1, 32),
106
+ 'max_model_len-4096-sw-None': (32, 16),
107
+ 'max_model_len-1024-sw-128': (2, 16),
108
+ 'max_model_len-512-sw-128': (2, 16),
109
+ 'max_model_len-256-sw-128': (2, 16),
110
+ 'max_model_len-2048-sw-None': (16, 32),
111
+ 'max_model_len-128-sw-128': (1, 32),
112
+ },
113
+ 'q_head-16_kv_head-2_head-64': {
114
+ 'max_model_len-128-sw-128': (1, 32),
115
+ 'max_model_len-256-sw-None': (2, 32),
116
+ 'max_model_len-256-sw-128': (2, 32),
117
+ 'max_model_len-512-sw-None': (4, 32),
118
+ 'max_model_len-512-sw-128': (2, 32),
119
+ 'max_model_len-1024-sw-None': (8, 32),
120
+ 'max_model_len-1024-sw-128': (2, 64),
121
+ 'max_model_len-2048-sw-None': (16, 32),
122
+ 'max_model_len-2048-sw-128': (2, 32),
123
+ 'max_model_len-4096-sw-None': (32, 32),
124
+ 'max_model_len-4096-sw-128': (2, 32),
125
+ 'max_model_len-8192-sw-None': (32, 32),
126
+ 'max_model_len-8192-sw-128': (2, 64),
127
+ 'max_model_len-128-sw-None': (1, 64),
128
+ },
129
+ 'q_head-8_kv_head-2_head-64': {
130
+ 'max_model_len-1024-sw-128': (2, 64),
131
+ 'max_model_len-2048-sw-None': (16, 32),
132
+ 'max_model_len-2048-sw-128': (2, 32),
133
+ 'max_model_len-4096-sw-None': (32, 32),
134
+ 'max_model_len-4096-sw-128': (2, 128),
135
+ 'max_model_len-8192-sw-None': (32, 32),
136
+ 'max_model_len-8192-sw-128': (2, 64),
137
+ 'max_model_len-128-sw-None': (1, 128),
138
+ 'max_model_len-128-sw-128': (1, 64),
139
+ 'max_model_len-256-sw-None': (2, 64),
140
+ 'max_model_len-256-sw-128': (2, 32),
141
+ 'max_model_len-512-sw-None': (4, 64),
142
+ 'max_model_len-512-sw-128': (2, 64),
143
+ 'max_model_len-1024-sw-None': (8, 32),
144
+ },
145
+ 'q_head-32_kv_head-4_head-64': {
146
+ 'max_model_len-128-sw-None': (1, 32),
147
+ 'max_model_len-128-sw-128': (1, 32),
148
+ 'max_model_len-256-sw-None': (2, 32),
149
+ 'max_model_len-256-sw-128': (2, 32),
150
+ 'max_model_len-2048-sw-128': (2, 32),
151
+ 'max_model_len-512-sw-None': (4, 32),
152
+ 'max_model_len-512-sw-128': (2, 32),
153
+ 'max_model_len-1024-sw-None': (8, 32),
154
+ 'max_model_len-1024-sw-128': (2, 32),
155
+ 'max_model_len-2048-sw-None': (16, 32),
156
+ 'max_model_len-4096-sw-None': (32, 32),
157
+ 'max_model_len-4096-sw-128': (2, 32),
158
+ 'max_model_len-8192-sw-None': (32, 32),
159
+ 'max_model_len-8192-sw-128': (2, 32),
160
+ },
161
+ }
162
+ },
163
+ },
164
+ 'TPU v6e': {
165
+ 256: {
166
+ 'q_bfloat16_kv_bfloat16': {
167
+ 'q_head-16_kv_head-2_head-64': {
168
+ 'max_model_len-2048-sw-None': (8, 32),
169
+ 'max_model_len-2048-sw-128': (1, 32),
170
+ 'max_model_len-4096-sw-None': (16, 32),
171
+ 'max_model_len-4096-sw-128': (1, 32),
172
+ 'max_model_len-8192-sw-None': (16, 32),
173
+ 'max_model_len-8192-sw-128': (1, 32),
174
+ 'max_model_len-256-sw-None': (1, 32),
175
+ 'max_model_len-256-sw-128': (1, 32),
176
+ 'max_model_len-512-sw-None': (2, 32),
177
+ 'max_model_len-512-sw-128': (1, 32),
178
+ 'max_model_len-1024-sw-None': (4, 32),
179
+ 'max_model_len-1024-sw-128': (1, 32),
180
+ },
181
+ 'q_head-64_kv_head-8_head-64': {
182
+ 'max_model_len-8192-sw-128': (1, 16),
183
+ 'max_model_len-512-sw-None': (2, 16),
184
+ 'max_model_len-256-sw-None': (1, 16),
185
+ 'max_model_len-256-sw-128': (1, 16),
186
+ 'max_model_len-512-sw-128': (1, 16),
187
+ 'max_model_len-1024-sw-None': (4, 16),
188
+ 'max_model_len-1024-sw-128': (1, 16),
189
+ 'max_model_len-2048-sw-None': (8, 32),
190
+ 'max_model_len-2048-sw-128': (1, 16),
191
+ 'max_model_len-4096-sw-None': (16, 16),
192
+ 'max_model_len-4096-sw-128': (1, 16),
193
+ 'max_model_len-8192-sw-None': (16, 16),
194
+ },
195
+ 'q_head-32_kv_head-4_head-64': {
196
+ 'max_model_len-256-sw-None': (1, 32),
197
+ 'max_model_len-256-sw-128': (1, 32),
198
+ 'max_model_len-512-sw-None': (2, 32),
199
+ 'max_model_len-512-sw-128': (1, 32),
200
+ 'max_model_len-4096-sw-None': (16, 32),
201
+ 'max_model_len-1024-sw-None': (4, 32),
202
+ 'max_model_len-1024-sw-128': (1, 32),
203
+ 'max_model_len-2048-sw-None': (8, 32),
204
+ 'max_model_len-2048-sw-128': (1, 32),
205
+ 'max_model_len-4096-sw-128': (1, 32),
206
+ 'max_model_len-8192-sw-None': (16, 32),
207
+ 'max_model_len-8192-sw-128': (1, 32),
208
+ },
209
+ 'q_head-8_kv_head-2_head-64': {
210
+ 'max_model_len-256-sw-None': (1, 64),
211
+ 'max_model_len-256-sw-128': (1, 64),
212
+ 'max_model_len-512-sw-None': (2, 64),
213
+ 'max_model_len-512-sw-128': (1, 64),
214
+ 'max_model_len-1024-sw-None': (4, 32),
215
+ 'max_model_len-1024-sw-128': (1, 64),
216
+ 'max_model_len-2048-sw-None': (8, 32),
217
+ 'max_model_len-2048-sw-128': (1, 64),
218
+ 'max_model_len-4096-sw-None': (16, 32),
219
+ 'max_model_len-4096-sw-128': (1, 32),
220
+ 'max_model_len-8192-sw-None': (16, 32),
221
+ 'max_model_len-8192-sw-128': (1, 32),
222
+ },
223
+ }
224
+ },
225
+ 128: {
226
+ 'q_bfloat16_kv_bfloat16': {
227
+ 'q_head-64_kv_head-8_head-64': {
228
+ 'max_model_len-4096-sw-128': (2, 16),
229
+ 'max_model_len-8192-sw-None': (16, 32),
230
+ 'max_model_len-1024-sw-None': (8, 16),
231
+ 'max_model_len-8192-sw-128': (2, 16),
232
+ 'max_model_len-512-sw-None': (4, 16),
233
+ 'max_model_len-2048-sw-128': (2, 16),
234
+ 'max_model_len-256-sw-None': (2, 16),
235
+ 'max_model_len-128-sw-None': (1, 32),
236
+ 'max_model_len-4096-sw-None': (32, 16),
237
+ 'max_model_len-1024-sw-128': (2, 16),
238
+ 'max_model_len-512-sw-128': (2, 16),
239
+ 'max_model_len-256-sw-128': (2, 16),
240
+ 'max_model_len-2048-sw-None': (16, 32),
241
+ 'max_model_len-128-sw-128': (1, 32),
242
+ },
243
+ 'q_head-16_kv_head-2_head-64': {
244
+ 'max_model_len-128-sw-128': (1, 32),
245
+ 'max_model_len-256-sw-None': (2, 32),
246
+ 'max_model_len-256-sw-128': (2, 32),
247
+ 'max_model_len-512-sw-None': (4, 32),
248
+ 'max_model_len-512-sw-128': (2, 32),
249
+ 'max_model_len-1024-sw-None': (8, 32),
250
+ 'max_model_len-1024-sw-128': (2, 32),
251
+ 'max_model_len-2048-sw-None': (16, 32),
252
+ 'max_model_len-2048-sw-128': (2, 32),
253
+ 'max_model_len-4096-sw-None': (32, 32),
254
+ 'max_model_len-4096-sw-128': (2, 32),
255
+ 'max_model_len-8192-sw-None': (32, 32),
256
+ 'max_model_len-8192-sw-128': (2, 32),
257
+ 'max_model_len-128-sw-None': (1, 32),
258
+ },
259
+ 'q_head-8_kv_head-2_head-64': {
260
+ 'max_model_len-1024-sw-128': (2, 32),
261
+ 'max_model_len-2048-sw-None': (16, 32),
262
+ 'max_model_len-2048-sw-128': (2, 64),
263
+ 'max_model_len-4096-sw-None': (32, 32),
264
+ 'max_model_len-4096-sw-128': (2, 32),
265
+ 'max_model_len-8192-sw-None': (32, 32),
266
+ 'max_model_len-8192-sw-128': (2, 64),
267
+ 'max_model_len-128-sw-None': (1, 64),
268
+ 'max_model_len-128-sw-128': (1, 64),
269
+ 'max_model_len-256-sw-None': (2, 32),
270
+ 'max_model_len-256-sw-128': (2, 32),
271
+ 'max_model_len-512-sw-None': (4, 32),
272
+ 'max_model_len-512-sw-128': (2, 64),
273
+ 'max_model_len-1024-sw-None': (8, 32),
274
+ },
275
+ 'q_head-32_kv_head-4_head-64': {
276
+ 'max_model_len-128-sw-None': (1, 32),
277
+ 'max_model_len-128-sw-128': (1, 32),
278
+ 'max_model_len-256-sw-None': (2, 32),
279
+ 'max_model_len-256-sw-128': (2, 32),
280
+ 'max_model_len-2048-sw-128': (2, 32),
281
+ 'max_model_len-512-sw-None': (4, 32),
282
+ 'max_model_len-512-sw-128': (2, 32),
283
+ 'max_model_len-1024-sw-None': (8, 32),
284
+ 'max_model_len-1024-sw-128': (2, 32),
285
+ 'max_model_len-2048-sw-None': (16, 32),
286
+ 'max_model_len-4096-sw-None': (32, 32),
287
+ 'max_model_len-4096-sw-128': (2, 32),
288
+ 'max_model_len-8192-sw-None': (32, 32),
289
+ 'max_model_len-8192-sw-128': (2, 32),
290
+ },
291
+ }
292
+ },
293
+ },
294
+ 'TPU v5e': {
295
+ 256: {
296
+ 'q_bfloat16_kv_bfloat16': {
297
+ 'q_head-16_kv_head-2_head-64': {
298
+ 'max_model_len-2048-sw-None': (8, 32),
299
+ 'max_model_len-2048-sw-128': (1, 32),
300
+ 'max_model_len-4096-sw-None': (16, 32),
301
+ 'max_model_len-4096-sw-128': (1, 32),
302
+ 'max_model_len-8192-sw-None': (16, 32),
303
+ 'max_model_len-8192-sw-128': (1, 32),
304
+ 'max_model_len-256-sw-None': (1, 32),
305
+ 'max_model_len-256-sw-128': (1, 32),
306
+ 'max_model_len-512-sw-None': (2, 32),
307
+ 'max_model_len-512-sw-128': (1, 32),
308
+ 'max_model_len-1024-sw-None': (4, 32),
309
+ 'max_model_len-1024-sw-128': (1, 32),
310
+ },
311
+ 'q_head-64_kv_head-8_head-64': {
312
+ 'max_model_len-8192-sw-128': (1, 8),
313
+ 'max_model_len-512-sw-None': (2, 16),
314
+ 'max_model_len-256-sw-None': (1, 8),
315
+ 'max_model_len-256-sw-128': (1, 8),
316
+ 'max_model_len-512-sw-128': (1, 8),
317
+ 'max_model_len-1024-sw-None': (4, 32),
318
+ 'max_model_len-1024-sw-128': (1, 8),
319
+ 'max_model_len-2048-sw-None': (4, 32),
320
+ 'max_model_len-2048-sw-128': (1, 8),
321
+ 'max_model_len-4096-sw-None': (4, 32),
322
+ 'max_model_len-4096-sw-128': (1, 8),
323
+ 'max_model_len-8192-sw-None': (4, 32),
324
+ },
325
+ 'q_head-32_kv_head-4_head-64': {
326
+ 'max_model_len-256-sw-None': (1, 16),
327
+ 'max_model_len-256-sw-128': (1, 16),
328
+ 'max_model_len-512-sw-None': (2, 16),
329
+ 'max_model_len-512-sw-128': (1, 16),
330
+ 'max_model_len-4096-sw-None': (8, 32),
331
+ 'max_model_len-1024-sw-None': (4, 16),
332
+ 'max_model_len-1024-sw-128': (1, 16),
333
+ 'max_model_len-2048-sw-None': (8, 32),
334
+ 'max_model_len-2048-sw-128': (1, 16),
335
+ 'max_model_len-4096-sw-128': (1, 16),
336
+ 'max_model_len-8192-sw-None': (8, 32),
337
+ 'max_model_len-8192-sw-128': (1, 16),
338
+ },
339
+ 'q_head-8_kv_head-2_head-64': {
340
+ 'max_model_len-256-sw-None': (1, 32),
341
+ 'max_model_len-256-sw-128': (1, 32),
342
+ 'max_model_len-512-sw-None': (2, 32),
343
+ 'max_model_len-512-sw-128': (1, 32),
344
+ 'max_model_len-1024-sw-None': (4, 32),
345
+ 'max_model_len-1024-sw-128': (1, 32),
346
+ 'max_model_len-2048-sw-None': (8, 32),
347
+ 'max_model_len-2048-sw-128': (1, 32),
348
+ 'max_model_len-4096-sw-None': (16, 32),
349
+ 'max_model_len-4096-sw-128': (1, 32),
350
+ 'max_model_len-8192-sw-None': (8, 32),
351
+ 'max_model_len-8192-sw-128': (1, 32),
352
+ },
353
+ }
354
+ },
355
+ 128: {
356
+ 'q_bfloat16_kv_bfloat16': {
357
+ 'q_head-64_kv_head-8_head-64': {
358
+ 'max_model_len-4096-sw-128': (2, 8),
359
+ 'max_model_len-8192-sw-None': (8, 32),
360
+ 'max_model_len-1024-sw-None': (8, 32),
361
+ 'max_model_len-8192-sw-128': (2, 8),
362
+ 'max_model_len-512-sw-None': (4, 16),
363
+ 'max_model_len-2048-sw-128': (2, 8),
364
+ 'max_model_len-256-sw-None': (2, 8),
365
+ 'max_model_len-128-sw-None': (1, 8),
366
+ 'max_model_len-4096-sw-None': (8, 32),
367
+ 'max_model_len-1024-sw-128': (2, 16),
368
+ 'max_model_len-512-sw-128': (2, 8),
369
+ 'max_model_len-256-sw-128': (2, 8),
370
+ 'max_model_len-2048-sw-None': (8, 32),
371
+ 'max_model_len-128-sw-128': (1, 8),
372
+ },
373
+ 'q_head-16_kv_head-2_head-64': {
374
+ 'max_model_len-128-sw-128': (1, 32),
375
+ 'max_model_len-256-sw-None': (2, 32),
376
+ 'max_model_len-256-sw-128': (2, 32),
377
+ 'max_model_len-512-sw-None': (4, 32),
378
+ 'max_model_len-512-sw-128': (2, 32),
379
+ 'max_model_len-1024-sw-None': (8, 32),
380
+ 'max_model_len-1024-sw-128': (2, 32),
381
+ 'max_model_len-2048-sw-None': (16, 32),
382
+ 'max_model_len-2048-sw-128': (2, 32),
383
+ 'max_model_len-4096-sw-None': (32, 32),
384
+ 'max_model_len-4096-sw-128': (2, 32),
385
+ 'max_model_len-8192-sw-None': (32, 32),
386
+ 'max_model_len-8192-sw-128': (2, 32),
387
+ 'max_model_len-128-sw-None': (1, 32),
388
+ },
389
+ 'q_head-8_kv_head-2_head-64': {
390
+ 'max_model_len-1024-sw-128': (2, 32),
391
+ 'max_model_len-2048-sw-None': (16, 32),
392
+ 'max_model_len-2048-sw-128': (2, 32),
393
+ 'max_model_len-4096-sw-None': (32, 32),
394
+ 'max_model_len-4096-sw-128': (2, 32),
395
+ 'max_model_len-8192-sw-None': (32, 32),
396
+ 'max_model_len-8192-sw-128': (2, 32),
397
+ 'max_model_len-128-sw-None': (1, 32),
398
+ 'max_model_len-128-sw-128': (1, 32),
399
+ 'max_model_len-256-sw-None': (2, 32),
400
+ 'max_model_len-256-sw-128': (2, 32),
401
+ 'max_model_len-512-sw-None': (4, 32),
402
+ 'max_model_len-512-sw-128': (2, 32),
403
+ 'max_model_len-1024-sw-None': (8, 32),
404
+ },
405
+ 'q_head-32_kv_head-4_head-64': {
406
+ 'max_model_len-128-sw-None': (1, 16),
407
+ 'max_model_len-128-sw-128': (1, 16),
408
+ 'max_model_len-256-sw-None': (2, 16),
409
+ 'max_model_len-256-sw-128': (2, 16),
410
+ 'max_model_len-2048-sw-128': (2, 16),
411
+ 'max_model_len-512-sw-None': (4, 16),
412
+ 'max_model_len-512-sw-128': (2, 16),
413
+ 'max_model_len-1024-sw-None': (8, 16),
414
+ 'max_model_len-1024-sw-128': (2, 16),
415
+ 'max_model_len-2048-sw-None': (16, 32),
416
+ 'max_model_len-4096-sw-None': (16, 32),
417
+ 'max_model_len-4096-sw-128': (2, 16),
418
+ 'max_model_len-8192-sw-None': (16, 32),
419
+ 'max_model_len-8192-sw-128': (2, 16),
420
+ },
421
+ }
422
+ },
423
+ },
424
+ }
425
+
426
+
427
+ def get_tuned_block_sizes(
428
+ q_dtype,
429
+ kv_dtype,
430
+ actual_num_q_heads,
431
+ actual_num_kv_heads,
432
+ head_dim,
433
+ page_size,
434
+ max_num_tokens,
435
+ pages_per_seq,
436
+ sliding_window=None,
437
+ ) -> tuple[int, int]:
438
+ """Search tuned values for (num_kv_pages_per_blk, num_queries_per_blk)."""
439
+
440
+ keys = get_lookup_keys(
441
+ page_size,
442
+ q_dtype,
443
+ kv_dtype,
444
+ actual_num_q_heads,
445
+ actual_num_kv_heads,
446
+ head_dim,
447
+ page_size * pages_per_seq,
448
+ sliding_window,
449
+ )
450
+ device, page_size, dtypes, head_dims, extra = keys
451
+
452
+ try:
453
+ bkv_p, bq = TUNED_BLOCK_SIZES[device][page_size][dtypes][head_dims][
454
+ extra]
455
+ except KeyError:
456
+ logger.warning_once(
457
+ 'Couldn`t find tuned sizes for the RPA v3 kernel with %s', keys)
458
+ # When not available use a sensible default based on TPU version
459
+ # Set default block sizes for each tpu_version.
460
+ tpu_version = get_tpu_version()
461
+ if tpu_version < 4:
462
+ raise NotImplementedError('TPU version must be 4 or higher.')
463
+ match tpu_version:
464
+ case 4:
465
+ # TPUv4 has much smaller VMEM size so we pick fixed block sizes.
466
+ bkv_p, bq = (512 // page_size, 32)
467
+ case 7:
468
+ bkv_p, bq = (4096 // page_size, 32)
469
+ case _:
470
+ bkv_p, bq = (2048 // page_size, 32)
471
+
472
+ bkv_p, bq = (min(pages_per_seq, bkv_p), min(max_num_tokens, bq))
473
+
474
+ logger.info_once('RPA v3 kernel tuned block sizes for %s: bkv_p=%s, bq=%s',
475
+ keys, bkv_p, bq)
476
+ return bkv_p, bq
477
+
478
+
479
+ def get_lookup_keys(
480
+ page_size,
481
+ q_dtype,
482
+ kv_dtype,
483
+ num_q_heads,
484
+ num_kv_heads,
485
+ head_dim,
486
+ max_model_len,
487
+ sliding_window,
488
+ ):
489
+ """Get the lookup keys for tuned block sizes."""
490
+ (
491
+ page_size,
492
+ q_dtype_name,
493
+ kv_dtype_name,
494
+ num_q_heads,
495
+ num_kv_heads,
496
+ head_dim,
497
+ max_model_len,
498
+ sliding_window,
499
+ ) = get_simplified_raw_key(
500
+ page_size,
501
+ q_dtype,
502
+ kv_dtype,
503
+ num_q_heads,
504
+ num_kv_heads,
505
+ head_dim,
506
+ max_model_len,
507
+ sliding_window,
508
+ )
509
+
510
+ return (
511
+ get_device_name(),
512
+ next_power_of_2(page_size),
513
+ f'q_{q_dtype_name}_kv_{kv_dtype_name}',
514
+ f'q_head-{num_q_heads}_kv_head-{num_kv_heads}_head-{head_dim}',
515
+ f'max_model_len-{next_power_of_2(max_model_len)}-sw-{sliding_window}',
516
+ )
517
+
518
+
519
+ def get_simplified_raw_key(
520
+ page_size,
521
+ q_dtype,
522
+ kv_dtype,
523
+ actual_num_q_heads,
524
+ actual_num_kv_heads,
525
+ head_dim,
526
+ max_model_len,
527
+ sliding_window,
528
+ ):
529
+ """Get the simplified key."""
530
+ assert head_dim == 64
531
+ assert actual_num_q_heads % actual_num_kv_heads == 0
532
+ actual_num_q_heads_per_kv_head = actual_num_q_heads // actual_num_kv_heads
533
+ q_packing = get_dtype_packing(q_dtype)
534
+ kv_packing = get_dtype_packing(kv_dtype)
535
+ num_kv_heads = align_to(actual_num_kv_heads, kv_packing)
536
+ num_q_heads_per_kv_head = align_to(actual_num_q_heads_per_kv_head,
537
+ q_packing)
538
+
539
+ return (
540
+ next_power_of_2(page_size),
541
+ jnp.dtype(q_dtype).name,
542
+ jnp.dtype(kv_dtype).name,
543
+ next_power_of_2(num_q_heads_per_kv_head * actual_num_kv_heads),
544
+ next_power_of_2(num_kv_heads),
545
+ head_dim,
546
+ next_power_of_2(max_model_len),
547
+ sliding_window,
548
+ )
tpu_inference/kernels/ragged_paged_attention/v3/util.py ADDED
@@ -0,0 +1,65 @@
1
+ # Copyright 2025 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """Utility functions for ragged paged attention."""
15
+ import jax
16
+ from jax._src import dtypes
17
+
18
+
19
+ def cdiv(a, b):
20
+ assert b != 0
21
+ return (a + b - 1) // b
22
+
23
+
24
+ def align_to(x, a):
25
+ return cdiv(x, a) * a
26
+
27
+
28
+ def get_dtype_bitwidth(dtype):
29
+ return (dtypes.bit_width(dtype)
30
+ if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
31
+
32
+
33
+ def get_dtype_packing(dtype):
34
+ bits = get_dtype_bitwidth(dtype)
35
+ return 32 // bits
36
+
37
+
38
+ def next_power_of_2(x: int):
39
+ """Finds the smallest power of 2 >= x using bit manipulation.
40
+
41
+ Args:
42
+ x: The input number (should be an integer).
43
+
44
+ Returns:
45
+ The smallest integer power of 2 that is >= x.
46
+ """
47
+ assert x > 0
48
+ if x == 1:
49
+ return 1
50
+ return 1 << (x - 1).bit_length()
51
+
52
+
53
+ def get_tpu_version() -> int:
54
+ """Returns the numeric version of the TPU, or -1 if not on TPU."""
55
+ kind = jax.devices()[0].device_kind
56
+ if 'TPU' not in kind:
57
+ return -1
58
+ if kind.endswith(' lite'):
59
+ kind = kind[:-len(' lite')]
60
+ if kind.endswith('p') or kind.endswith('e'):
61
+ kind = kind[:-1]
62
+ if kind == 'TPU7x':
63
+ return 7
64
+ assert kind[:-1] == 'TPU v', kind
65
+ return int(kind[-1])
tpu_inference/layers/__init__.py ADDED
@@ -0,0 +1,13 @@
1
+ # Copyright 2025 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
tpu_inference/layers/common/__init__.py ADDED
@@ -0,0 +1,13 @@
1
+ # Copyright 2025 Google LLC
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.