tpu-inference 0.12.0.dev20251217__tar.gz → 0.13.2.dev20260104__tar.gz

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.

Potentially problematic release.


This version of tpu-inference might be problematic. Click here for more details.

Files changed (332) hide show
  1. {tpu_inference-0.12.0.dev20251217/tpu_inference.egg-info → tpu_inference-0.13.2.dev20260104}/PKG-INFO +1 -1
  2. tpu_inference-0.13.2.dev20260104/pyproject.toml +13 -0
  3. tpu_inference-0.13.2.dev20260104/tests/__init__.py +13 -0
  4. tpu_inference-0.13.2.dev20260104/tests/core/__init__.py +13 -0
  5. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/core/test_disagg_utils.py +14 -0
  6. tpu_inference-0.13.2.dev20260104/tests/core/test_dp_scheduler.py +781 -0
  7. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/core/test_init.py +14 -0
  8. tpu_inference-0.13.2.dev20260104/tests/distributed/__init__.py +13 -0
  9. tpu_inference-0.13.2.dev20260104/tests/distributed/test_distributed_utils.py +120 -0
  10. tpu_inference-0.13.2.dev20260104/tests/distributed/test_tpu_connector.py +478 -0
  11. tpu_inference-0.13.2.dev20260104/tests/e2e/__init__.py +13 -0
  12. tpu_inference-0.13.2.dev20260104/tests/e2e/test_async_scheduler.py +211 -0
  13. tpu_inference-0.13.2.dev20260104/tests/e2e/test_data_parallel.py +289 -0
  14. tpu_inference-0.13.2.dev20260104/tests/e2e/test_hybrid_kvcache.py +219 -0
  15. tpu_inference-0.13.2.dev20260104/tests/e2e/test_local_disagg.py +257 -0
  16. tpu_inference-0.13.2.dev20260104/tests/e2e/test_model_loader.py +268 -0
  17. tpu_inference-0.13.2.dev20260104/tests/e2e/test_multi_modal_inference.py +111 -0
  18. tpu_inference-0.13.2.dev20260104/tests/e2e/test_pipeline_parallel.py +265 -0
  19. tpu_inference-0.13.2.dev20260104/tests/e2e/test_runai_model_streamer_loader.py +104 -0
  20. tpu_inference-0.13.2.dev20260104/tests/e2e/test_sampling_params.py +269 -0
  21. tpu_inference-0.13.2.dev20260104/tests/e2e/test_speculative_decoding.py +311 -0
  22. tpu_inference-0.13.2.dev20260104/tests/e2e/test_structured_decoding.py +46 -0
  23. tpu_inference-0.13.2.dev20260104/tests/executors/__init__.py +13 -0
  24. tpu_inference-0.13.2.dev20260104/tests/executors/test_ray_distributed_executor.py +199 -0
  25. tpu_inference-0.13.2.dev20260104/tests/experimental/__init__.py +13 -0
  26. tpu_inference-0.13.2.dev20260104/tests/experimental/test_llama3_jax_stashed.py +208 -0
  27. tpu_inference-0.13.2.dev20260104/tests/kernels/__init__.py +13 -0
  28. tpu_inference-0.13.2.dev20260104/tests/kernels/collectives/__init__.py +13 -0
  29. tpu_inference-0.13.2.dev20260104/tests/kernels/collectives/all_gather_matmul_kernel_test.py +69 -0
  30. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/fused_moe_v1_test.py +14 -0
  31. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/gmm_test.py +14 -0
  32. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/mla_v1_test.py +14 -0
  33. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/ragged_kv_cache_update_v2_test.py +14 -0
  34. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/ragged_paged_attention_kernel_v2_test.py +14 -0
  35. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py +14 -0
  36. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/ragged_paged_attention_kernel_v3_test.py +14 -0
  37. tpu_inference-0.13.2.dev20260104/tests/layers/__init__.py +13 -0
  38. tpu_inference-0.13.2.dev20260104/tests/layers/common/__init__.py +13 -0
  39. tpu_inference-0.13.2.dev20260104/tests/layers/common/test_attention_interface.py +156 -0
  40. tpu_inference-0.13.2.dev20260104/tests/layers/common/test_quantization.py +149 -0
  41. tpu_inference-0.13.2.dev20260104/tests/layers/jax/__init__.py +13 -0
  42. tpu_inference-0.13.2.dev20260104/tests/layers/jax/attention/__init__.py +13 -0
  43. tpu_inference-0.13.2.dev20260104/tests/layers/jax/attention/test_common_attention.py +103 -0
  44. tpu_inference-0.13.2.dev20260104/tests/layers/jax/attention/test_deepseek_v3_attention.py +233 -0
  45. tpu_inference-0.13.2.dev20260104/tests/layers/jax/attention/test_llama4_attention.py +135 -0
  46. tpu_inference-0.13.2.dev20260104/tests/layers/jax/moe/__init__.py +13 -0
  47. tpu_inference-0.13.2.dev20260104/tests/layers/jax/moe/test_deepseek_moe.py +235 -0
  48. tpu_inference-0.13.2.dev20260104/tests/layers/jax/sample/__init__.py +13 -0
  49. tpu_inference-0.13.2.dev20260104/tests/layers/jax/sample/test_rejection_sampler.py +1624 -0
  50. tpu_inference-0.13.2.dev20260104/tests/layers/jax/sample/test_sampling.py +115 -0
  51. tpu_inference-0.13.2.dev20260104/tests/layers/jax/sample/test_sampling_metadata.py +254 -0
  52. tpu_inference-0.13.2.dev20260104/tests/layers/jax/test_layers.py +155 -0
  53. tpu_inference-0.12.0.dev20251217/tests/test_quantization.py → tpu_inference-0.13.2.dev20260104/tests/layers/jax/test_qwix.py +180 -50
  54. tpu_inference-0.13.2.dev20260104/tests/layers/jax/test_rope.py +93 -0
  55. tpu_inference-0.13.2.dev20260104/tests/layers/jax/test_sharding.py +159 -0
  56. tpu_inference-0.13.2.dev20260104/tests/layers/jax/test_transformer_block.py +152 -0
  57. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/__init__.py +13 -0
  58. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_attention.py +363 -0
  59. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_awq.py +406 -0
  60. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_compressed_tensors_moe.py +199 -0
  61. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_compressed_tensors_w8a8_fp8.py +441 -0
  62. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_compressed_tensors_w8a8_int8.py +443 -0
  63. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_fp8.py +17 -0
  64. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_mxfp4.py +320 -0
  65. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/test_unquantized.py +662 -0
  66. tpu_inference-0.13.2.dev20260104/tests/layers/vllm/utils.py +87 -0
  67. tpu_inference-0.13.2.dev20260104/tests/lora/__init__.py +13 -0
  68. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/conftest.py +14 -0
  69. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/test_bgmv.py +14 -0
  70. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/test_layers.py +19 -3
  71. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/test_lora.py +14 -0
  72. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/test_lora_perf.py +14 -0
  73. tpu_inference-0.13.2.dev20260104/tests/models/__init__.py +13 -0
  74. tpu_inference-0.13.2.dev20260104/tests/models/common/__init__.py +13 -0
  75. tpu_inference-0.13.2.dev20260104/tests/models/common/test_model_loader.py +455 -0
  76. tpu_inference-0.13.2.dev20260104/tests/models/jax/__init__.py +13 -0
  77. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_deepseek_v3.py +401 -0
  78. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_llama3.py +184 -0
  79. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_llama4.py +298 -0
  80. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_llama_eagle3.py +197 -0
  81. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_llama_guard_4.py +242 -0
  82. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_qwen2.py +172 -0
  83. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_qwen2_5_vl.py +605 -0
  84. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_qwen3.py +169 -0
  85. tpu_inference-0.13.2.dev20260104/tests/models/jax/test_weight_loading.py +180 -0
  86. tpu_inference-0.13.2.dev20260104/tests/models/jax/utils/__init__.py +13 -0
  87. tpu_inference-0.13.2.dev20260104/tests/models/jax/utils/test_multi_modal_utils.py +212 -0
  88. tpu_inference-0.13.2.dev20260104/tests/platforms/__init__.py +13 -0
  89. tpu_inference-0.13.2.dev20260104/tests/platforms/test_tpu_platform.py +54 -0
  90. tpu_inference-0.13.2.dev20260104/tests/runner/__init__.py +13 -0
  91. tpu_inference-0.13.2.dev20260104/tests/runner/test_block_table.py +395 -0
  92. tpu_inference-0.13.2.dev20260104/tests/runner/test_input_batch.py +226 -0
  93. tpu_inference-0.13.2.dev20260104/tests/runner/test_kv_cache.py +220 -0
  94. tpu_inference-0.13.2.dev20260104/tests/runner/test_kv_cache_manager.py +498 -0
  95. tpu_inference-0.13.2.dev20260104/tests/runner/test_multimodal_manager.py +429 -0
  96. tpu_inference-0.13.2.dev20260104/tests/runner/test_persistent_batch_manager.py +84 -0
  97. tpu_inference-0.13.2.dev20260104/tests/runner/test_speculative_decoding_manager.py +368 -0
  98. tpu_inference-0.13.2.dev20260104/tests/runner/test_structured_decoding_manager.py +220 -0
  99. tpu_inference-0.13.2.dev20260104/tests/runner/test_tpu_runner.py +261 -0
  100. tpu_inference-0.13.2.dev20260104/tests/runner/test_tpu_runner_dp.py +1099 -0
  101. tpu_inference-0.13.2.dev20260104/tests/runner/test_tpu_runner_mesh.py +200 -0
  102. tpu_inference-0.13.2.dev20260104/tests/runner/test_utils.py +411 -0
  103. tpu_inference-0.13.2.dev20260104/tests/spec_decode/__init__.py +13 -0
  104. tpu_inference-0.13.2.dev20260104/tests/spec_decode/test_eagle3.py +311 -0
  105. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/test_base.py +14 -0
  106. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/test_tpu_info.py +14 -0
  107. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/test_utils.py +1 -43
  108. tpu_inference-0.13.2.dev20260104/tests/worker/__init__.py +13 -0
  109. tpu_inference-0.13.2.dev20260104/tests/worker/tpu_worker_test.py +414 -0
  110. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/__init__.py +14 -0
  111. tpu_inference-0.13.2.dev20260104/tpu_inference/core/__init__.py +13 -0
  112. tpu_inference-0.13.2.dev20260104/tpu_inference/core/sched/__init__.py +13 -0
  113. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/core/sched/dp_scheduler.py +372 -56
  114. tpu_inference-0.13.2.dev20260104/tpu_inference/distributed/__init__.py +13 -0
  115. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/distributed/jax_parallel_state.py +14 -0
  116. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/distributed/tpu_connector.py +11 -6
  117. tpu_inference-0.13.2.dev20260104/tpu_inference/distributed/utils.py +112 -0
  118. tpu_inference-0.13.2.dev20260104/tpu_inference/executors/__init__.py +13 -0
  119. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/executors/ray_distributed_executor.py +17 -3
  120. tpu_inference-0.13.2.dev20260104/tpu_inference/experimental/__init__.py +13 -0
  121. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/experimental/llama3_jax_stashed.py +14 -0
  122. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/__init__.py +13 -0
  123. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/collectives/__init__.py +13 -0
  124. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/flash_attention/__init__.py +13 -0
  125. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/fused_moe/__init__.py +13 -0
  126. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/fused_moe/v1/__init__.py +13 -0
  127. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/fused_moe/v1/kernel.py +13 -0
  128. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/megablox/__init__.py +13 -0
  129. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/megablox/common.py +13 -0
  130. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/megablox/gmm.py +13 -0
  131. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/mla/__init__.py +13 -0
  132. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/mla/v1/__init__.py +13 -0
  133. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/mla/v1/kernel.py +20 -26
  134. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/quantized_matmul/__init__.py +13 -0
  135. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/ragged_paged_attention/__init__.py +13 -0
  136. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/ragged_paged_attention/v2/__init__.py +13 -0
  137. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/ragged_paged_attention/v3/__init__.py +13 -0
  138. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v3/kernel.py +19 -1
  139. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py +19 -1
  140. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py +4460 -0
  141. tpu_inference-0.13.2.dev20260104/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py +548 -0
  142. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v3/util.py +13 -0
  143. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/__init__.py +13 -0
  144. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/common/__init__.py +13 -0
  145. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/common/attention_interface.py +25 -12
  146. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/common/attention_metadata.py +14 -0
  147. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/fused_moe.py → tpu_inference-0.13.2.dev20260104/tpu_inference/layers/common/fused_moe_gmm.py +72 -55
  148. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/common/quant_methods.py +23 -0
  149. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/common/quantization.py +282 -0
  150. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/common/sharding.py +22 -3
  151. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/common/utils.py +94 -0
  152. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/jax/__init__.py +13 -0
  153. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/jax/attention/__init__.py +13 -0
  154. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/attention/attention.py +19 -6
  155. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/attention/deepseek_v3_attention.py +52 -27
  156. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/attention/gpt_oss_attention.py +19 -6
  157. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/attention/llama4_attention.py +17 -4
  158. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/base.py +14 -0
  159. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/constants.py +13 -0
  160. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/layers.py +14 -0
  161. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/jax/misc.py +30 -0
  162. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/jax/moe/__init__.py +13 -0
  163. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/moe/deepseek_v3_moe.py +20 -13
  164. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/moe/gpt_oss_moe.py +14 -0
  165. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/moe/moe.py +43 -3
  166. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/pp_utils.py +14 -0
  167. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/rope.py +14 -0
  168. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/rope_interface.py +14 -0
  169. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/jax/sample/__init__.py +13 -0
  170. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/sample/rejection_sampler.py +13 -0
  171. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/sample/sampling.py +15 -1
  172. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/sample/sampling_metadata.py +14 -0
  173. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/jax/transformer_block.py +14 -0
  174. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/__init__.py +13 -0
  175. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/attention.py +4 -4
  176. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/fused_moe.py +114 -0
  177. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/linear.py +64 -0
  178. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/process_weights/__init__.py +13 -0
  179. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/sharding.py → tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/process_weights/cleanup_sharding.py +24 -15
  180. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/process_weights/fused_moe_weights.py +369 -0
  181. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/process_weights/linear_weights.py +174 -0
  182. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/quantization/__init__.py +19 -3
  183. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/quantization/awq.py +96 -82
  184. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py +13 -0
  185. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py +19 -5
  186. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py +199 -0
  187. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py +13 -0
  188. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +111 -91
  189. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +79 -43
  190. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/common.py → tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/configs.py +38 -26
  191. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/fp8.py +119 -0
  192. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/mxfp4.py +225 -0
  193. tpu_inference-0.13.2.dev20260104/tpu_inference/layers/vllm/quantization/unquantized.py +298 -0
  194. tpu_inference-0.13.2.dev20260104/tpu_inference/lora/__init__.py +13 -0
  195. tpu_inference-0.13.2.dev20260104/tpu_inference/models/__init__.py +13 -0
  196. tpu_inference-0.13.2.dev20260104/tpu_inference/models/common/__init__.py +13 -0
  197. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/common/model_loader.py +32 -16
  198. tpu_inference-0.13.2.dev20260104/tpu_inference/models/jax/__init__.py +13 -0
  199. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/deepseek_v3.py +113 -124
  200. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/gpt_oss.py +23 -7
  201. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/jax_intermediate_tensor.py +14 -0
  202. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/llama3.py +28 -11
  203. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/llama4.py +14 -0
  204. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/llama_eagle3.py +14 -0
  205. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/llama_guard_4.py +15 -1
  206. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/qwen2.py +17 -2
  207. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/qwen2_5_vl.py +18 -4
  208. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/qwen3.py +17 -2
  209. tpu_inference-0.13.2.dev20260104/tpu_inference/models/jax/utils/__init__.py +13 -0
  210. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/utils/file_utils.py +14 -0
  211. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/utils/multi_modal_utils.py +18 -4
  212. tpu_inference-0.13.2.dev20260104/tpu_inference/models/jax/utils/qwix/__init__.py +13 -0
  213. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/quantization_utils.py → tpu_inference-0.13.2.dev20260104/tpu_inference/models/jax/utils/qwix/qwix_utils.py +85 -24
  214. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/jax/utils/weight_utils.py +13 -0
  215. tpu_inference-0.13.2.dev20260104/tpu_inference/models/vllm/__init__.py +13 -0
  216. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/vllm/vllm_model_wrapper.py +21 -3
  217. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/models/vllm/vllm_model_wrapper_context.py +14 -0
  218. tpu_inference-0.13.2.dev20260104/tpu_inference/platforms/__init__.py +16 -0
  219. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/platforms/tpu_platform.py +22 -24
  220. tpu_inference-0.13.2.dev20260104/tpu_inference/runner/__init__.py +13 -0
  221. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/compilation_manager.py +79 -39
  222. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/kv_cache.py +14 -0
  223. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/kv_cache_manager.py +15 -2
  224. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/lora_utils.py +16 -1
  225. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/multimodal_manager.py +16 -2
  226. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/persistent_batch_manager.py +14 -0
  227. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/speculative_decoding_manager.py +14 -0
  228. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/structured_decoding_manager.py +14 -0
  229. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/tpu_runner.py +30 -10
  230. tpu_inference-0.13.2.dev20260104/tpu_inference/spec_decode/__init__.py +13 -0
  231. tpu_inference-0.13.2.dev20260104/tpu_inference/spec_decode/jax/__init__.py +13 -0
  232. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/spec_decode/jax/eagle3.py +13 -0
  233. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/tpu_info.py +14 -0
  234. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/utils.py +31 -30
  235. tpu_inference-0.13.2.dev20260104/tpu_inference/worker/__init__.py +13 -0
  236. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/worker/tpu_worker.py +19 -6
  237. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104/tpu_inference.egg-info}/PKG-INFO +1 -1
  238. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference.egg-info/SOURCES.txt +97 -11
  239. tpu_inference-0.12.0.dev20251217/pyproject.toml +0 -0
  240. tpu_inference-0.12.0.dev20251217/tests/__init__.py +0 -0
  241. tpu_inference-0.12.0.dev20251217/tests/core/__init__.py +0 -0
  242. tpu_inference-0.12.0.dev20251217/tests/core/test_dp_scheduler.py +0 -899
  243. tpu_inference-0.12.0.dev20251217/tests/kernels/__init__.py +0 -0
  244. tpu_inference-0.12.0.dev20251217/tests/lora/__init__.py +0 -0
  245. tpu_inference-0.12.0.dev20251217/tpu_inference/core/__init__.py +0 -0
  246. tpu_inference-0.12.0.dev20251217/tpu_inference/core/sched/__init__.py +0 -0
  247. tpu_inference-0.12.0.dev20251217/tpu_inference/distributed/__init__.py +0 -0
  248. tpu_inference-0.12.0.dev20251217/tpu_inference/distributed/utils.py +0 -60
  249. tpu_inference-0.12.0.dev20251217/tpu_inference/executors/__init__.py +0 -0
  250. tpu_inference-0.12.0.dev20251217/tpu_inference/experimental/__init__.py +0 -0
  251. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/__init__.py +0 -0
  252. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/collectives/__init__.py +0 -0
  253. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/flash_attention/__init__.py +0 -0
  254. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/fused_moe/__init__.py +0 -0
  255. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/fused_moe/v1/__init__.py +0 -0
  256. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/megablox/__init__.py +0 -0
  257. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/mla/__init__.py +0 -0
  258. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/mla/v1/__init__.py +0 -0
  259. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/quantized_matmul/__init__.py +0 -0
  260. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/__init__.py +0 -0
  261. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v2/__init__.py +0 -0
  262. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/__init__.py +0 -0
  263. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py +0 -4147
  264. tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py +0 -368
  265. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/__init__.py +0 -0
  266. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/common/__init__.py +0 -0
  267. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/common/quant_methods.py +0 -8
  268. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/__init__.py +0 -0
  269. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/attention/__init__.py +0 -0
  270. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/misc.py +0 -16
  271. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/moe/__init__.py +0 -0
  272. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/sample/__init__.py +0 -0
  273. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/__init__.py +0 -0
  274. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/linear_common.py +0 -208
  275. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py +0 -0
  276. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py +0 -252
  277. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py +0 -0
  278. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/mxfp4.py +0 -448
  279. tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/unquantized.py +0 -404
  280. tpu_inference-0.12.0.dev20251217/tpu_inference/lora/__init__.py +0 -0
  281. tpu_inference-0.12.0.dev20251217/tpu_inference/models/__init__.py +0 -0
  282. tpu_inference-0.12.0.dev20251217/tpu_inference/models/common/__init__.py +0 -0
  283. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/__init__.py +0 -0
  284. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/__init__.py +0 -0
  285. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/__init__.py +0 -0
  286. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/fp8_all_modules_w_only.yaml +0 -5
  287. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/fp8_default.yaml +0 -6
  288. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/int8_all_modules_w_only.yaml +0 -5
  289. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/int8_default.yaml +0 -6
  290. tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/mxfp4_utils.py +0 -105
  291. tpu_inference-0.12.0.dev20251217/tpu_inference/models/vllm/__init__.py +0 -0
  292. tpu_inference-0.12.0.dev20251217/tpu_inference/platforms/__init__.py +0 -2
  293. tpu_inference-0.12.0.dev20251217/tpu_inference/runner/__init__.py +0 -0
  294. tpu_inference-0.12.0.dev20251217/tpu_inference/spec_decode/__init__.py +0 -0
  295. tpu_inference-0.12.0.dev20251217/tpu_inference/spec_decode/jax/__init__.py +0 -0
  296. tpu_inference-0.12.0.dev20251217/tpu_inference/worker/__init__.py +0 -0
  297. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/LICENSE +0 -0
  298. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/MANIFEST.in +0 -0
  299. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/README.md +0 -0
  300. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/requirements.txt +0 -0
  301. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/setup.cfg +0 -0
  302. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/setup.py +0 -0
  303. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/core/test_core_tpu.py +0 -0
  304. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/core/test_disagg_executor.py +0 -0
  305. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/kernels/quantized_matmul_kernel_test.py +0 -0
  306. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/lora/utils.py +0 -0
  307. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tests/test_envs.py +0 -0
  308. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/core/core_tpu.py +0 -0
  309. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/core/disagg_executor.py +0 -0
  310. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/core/disagg_utils.py +0 -0
  311. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/env_override.py +0 -0
  312. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/envs.py +0 -0
  313. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/collectives/all_gather_matmul.py +0 -0
  314. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py +0 -0
  315. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/collectives/util.py +0 -0
  316. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/flash_attention/kernel.py +0 -0
  317. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/quantized_matmul/kernel.py +0 -0
  318. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py +0 -0
  319. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/quantized_matmul/util.py +0 -0
  320. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v2/kernel.py +0 -0
  321. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py +0 -0
  322. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py +0 -0
  323. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/layers/common/binary_search.py +0 -0
  324. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/logger.py +0 -0
  325. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/lora/torch_lora_ops.py +0 -0
  326. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/lora/torch_punica_tpu.py +0 -0
  327. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/block_table.py +0 -0
  328. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/input_batch.py +0 -0
  329. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference/runner/utils.py +0 -0
  330. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference.egg-info/dependency_links.txt +0 -0
  331. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference.egg-info/requires.txt +0 -0
  332. {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2.dev20260104}/tpu_inference.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tpu_inference
3
- Version: 0.12.0.dev20251217
3
+ Version: 0.13.2.dev20260104
4
4
  Author: tpu_inference Contributors
5
5
  Classifier: Development Status :: 3 - Alpha
6
6
  Classifier: Intended Audience :: Developers
@@ -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.
@@ -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.
@@ -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.
@@ -1,3 +1,17 @@
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
+
1
15
  import unittest
2
16
 
3
17
  from tpu_inference.core.disagg_utils import _parse_slices