tpu-inference 0.12.0.dev20251217__tar.gz → 0.13.2rc1__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.
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/MANIFEST.in +1 -0
- {tpu_inference-0.12.0.dev20251217/tpu_inference.egg-info → tpu_inference-0.13.2rc1}/PKG-INFO +1 -1
- tpu_inference-0.13.2rc1/pyproject.toml +13 -0
- tpu_inference-0.13.2rc1/requirements_v7x.txt +25 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/setup.py +19 -5
- tpu_inference-0.13.2rc1/tests/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/core/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/core/test_disagg_utils.py +14 -0
- tpu_inference-0.13.2rc1/tests/core/test_dp_scheduler.py +781 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/core/test_init.py +14 -0
- tpu_inference-0.13.2rc1/tests/distributed/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/distributed/test_distributed_utils.py +120 -0
- tpu_inference-0.13.2rc1/tests/distributed/test_tpu_connector.py +478 -0
- tpu_inference-0.13.2rc1/tests/e2e/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_async_scheduler.py +211 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_data_parallel.py +289 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_hybrid_kvcache.py +219 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_local_disagg.py +257 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_model_loader.py +268 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_multi_modal_inference.py +111 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_pipeline_parallel.py +265 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_runai_model_streamer_loader.py +104 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_sampling_params.py +269 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_speculative_decoding.py +311 -0
- tpu_inference-0.13.2rc1/tests/e2e/test_structured_decoding.py +46 -0
- tpu_inference-0.13.2rc1/tests/executors/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/executors/test_ray_distributed_executor.py +199 -0
- tpu_inference-0.13.2rc1/tests/experimental/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/experimental/test_llama3_jax_stashed.py +208 -0
- tpu_inference-0.13.2rc1/tests/kernels/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/kernels/collectives/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/kernels/collectives/all_gather_matmul_kernel_test.py +69 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/fused_moe_v1_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/gmm_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/mla_v1_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/ragged_kv_cache_update_v2_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/ragged_paged_attention_kernel_v2_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/ragged_paged_attention_kernel_v3_test.py +14 -0
- tpu_inference-0.13.2rc1/tests/layers/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/common/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/common/test_attention_interface.py +156 -0
- tpu_inference-0.13.2rc1/tests/layers/common/test_quantization.py +149 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/attention/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/attention/test_common_attention.py +103 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/attention/test_deepseek_v3_attention.py +233 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/attention/test_llama4_attention.py +135 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/moe/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/moe/test_deepseek_moe.py +235 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/sample/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/sample/test_rejection_sampler.py +1624 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/sample/test_sampling.py +115 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/sample/test_sampling_metadata.py +254 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/test_layers.py +155 -0
- tpu_inference-0.12.0.dev20251217/tests/test_quantization.py → tpu_inference-0.13.2rc1/tests/layers/jax/test_qwix.py +180 -50
- tpu_inference-0.13.2rc1/tests/layers/jax/test_rope.py +93 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/test_sharding.py +159 -0
- tpu_inference-0.13.2rc1/tests/layers/jax/test_transformer_block.py +152 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_attention.py +363 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_awq.py +405 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_compressed_tensors_moe.py +202 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_compressed_tensors_w8a8_fp8.py +418 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_compressed_tensors_w8a8_int8.py +441 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_fp8.py +17 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_mxfp4.py +312 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/test_unquantized.py +651 -0
- tpu_inference-0.13.2rc1/tests/layers/vllm/utils.py +87 -0
- tpu_inference-0.13.2rc1/tests/lora/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/conftest.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/test_bgmv.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/test_layers.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/test_lora.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/test_lora_perf.py +14 -0
- tpu_inference-0.13.2rc1/tests/models/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/models/common/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/models/common/test_model_loader.py +455 -0
- tpu_inference-0.13.2rc1/tests/models/jax/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_deepseek_v3.py +401 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_llama3.py +184 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_llama4.py +298 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_llama_eagle3.py +197 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_llama_guard_4.py +242 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_qwen2.py +172 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_qwen2_5_vl.py +605 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_qwen3.py +169 -0
- tpu_inference-0.13.2rc1/tests/models/jax/test_weight_loading.py +180 -0
- tpu_inference-0.13.2rc1/tests/models/jax/utils/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/models/jax/utils/test_multi_modal_utils.py +212 -0
- tpu_inference-0.13.2rc1/tests/platforms/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/platforms/test_tpu_platform.py +54 -0
- tpu_inference-0.13.2rc1/tests/runner/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/runner/test_block_table.py +395 -0
- tpu_inference-0.13.2rc1/tests/runner/test_input_batch.py +226 -0
- tpu_inference-0.13.2rc1/tests/runner/test_kv_cache.py +220 -0
- tpu_inference-0.13.2rc1/tests/runner/test_kv_cache_manager.py +498 -0
- tpu_inference-0.13.2rc1/tests/runner/test_multimodal_manager.py +429 -0
- tpu_inference-0.13.2rc1/tests/runner/test_persistent_batch_manager.py +84 -0
- tpu_inference-0.13.2rc1/tests/runner/test_speculative_decoding_manager.py +368 -0
- tpu_inference-0.13.2rc1/tests/runner/test_structured_decoding_manager.py +220 -0
- tpu_inference-0.13.2rc1/tests/runner/test_tpu_runner.py +261 -0
- tpu_inference-0.13.2rc1/tests/runner/test_tpu_runner_dp.py +1099 -0
- tpu_inference-0.13.2rc1/tests/runner/test_tpu_runner_mesh.py +200 -0
- tpu_inference-0.13.2rc1/tests/runner/test_utils.py +411 -0
- tpu_inference-0.13.2rc1/tests/spec_decode/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/spec_decode/test_eagle3.py +311 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/test_base.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/test_tpu_info.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/test_utils.py +1 -43
- tpu_inference-0.13.2rc1/tests/worker/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tests/worker/tpu_worker_test.py +414 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/__init__.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/core/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/core/sched/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/core/sched/dp_scheduler.py +372 -56
- tpu_inference-0.13.2rc1/tpu_inference/distributed/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/distributed/jax_parallel_state.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/distributed/tpu_connector.py +11 -6
- tpu_inference-0.13.2rc1/tpu_inference/distributed/utils.py +112 -0
- tpu_inference-0.13.2rc1/tpu_inference/executors/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/executors/ray_distributed_executor.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/experimental/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/experimental/llama3_jax_stashed.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/collectives/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/flash_attention/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/fused_moe/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/fused_moe/v1/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/fused_moe/v1/kernel.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/megablox/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/megablox/common.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/megablox/gmm.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/mla/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/mla/v1/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/mla/v1/kernel.py +20 -26
- tpu_inference-0.13.2rc1/tpu_inference/kernels/quantized_matmul/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/ragged_paged_attention/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/ragged_paged_attention/v2/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/ragged_paged_attention/v3/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v3/kernel.py +19 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py +19 -1
- tpu_inference-0.13.2rc1/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py +4460 -0
- tpu_inference-0.13.2rc1/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py +548 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v3/util.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/common/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/common/attention_interface.py +25 -12
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/common/attention_metadata.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/common/quant_methods.py +23 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/common/quantization.py +270 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/common/sharding.py +22 -3
- tpu_inference-0.13.2rc1/tpu_inference/layers/jax/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/jax/attention/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/attention/attention.py +19 -6
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/attention/deepseek_v3_attention.py +52 -27
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/attention/gpt_oss_attention.py +19 -6
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/attention/llama4_attention.py +17 -4
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/base.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/constants.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/layers.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/jax/misc.py +30 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/jax/moe/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/moe/deepseek_v3_moe.py +20 -13
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/moe/gpt_oss_moe.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/moe/moe.py +43 -3
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/pp_utils.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/rope.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/rope_interface.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/jax/sample/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/sample/rejection_sampler.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/sample/sampling.py +15 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/sample/sampling_metadata.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/jax/transformer_block.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/vllm/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/attention.py +4 -4
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/fused_moe.py +71 -54
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/linear_common.py +20 -7
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/__init__.py +16 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/awq.py +15 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/common.py +28 -13
- tpu_inference-0.13.2rc1/tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/layers/vllm/quantization/fp8.py +118 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/mxfp4.py +49 -87
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/quantization/unquantized.py +35 -11
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/vllm/sharding.py +21 -4
- tpu_inference-0.13.2rc1/tpu_inference/lora/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/models/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/models/common/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/common/model_loader.py +26 -15
- tpu_inference-0.13.2rc1/tpu_inference/models/jax/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/deepseek_v3.py +113 -124
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/gpt_oss.py +23 -7
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/jax_intermediate_tensor.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/llama3.py +28 -11
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/llama4.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/llama_eagle3.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/llama_guard_4.py +15 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/qwen2.py +17 -2
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/qwen2_5_vl.py +18 -4
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/qwen3.py +17 -2
- tpu_inference-0.13.2rc1/tpu_inference/models/jax/utils/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/utils/file_utils.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/utils/multi_modal_utils.py +18 -4
- tpu_inference-0.13.2rc1/tpu_inference/models/jax/utils/qwix/__init__.py +13 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/quantization_utils.py → tpu_inference-0.13.2rc1/tpu_inference/models/jax/utils/qwix/qwix_utils.py +84 -23
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/jax/utils/weight_utils.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/models/vllm/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/vllm/vllm_model_wrapper.py +19 -2
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/models/vllm/vllm_model_wrapper_context.py +14 -0
- tpu_inference-0.13.2rc1/tpu_inference/platforms/__init__.py +16 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/platforms/tpu_platform.py +28 -23
- tpu_inference-0.13.2rc1/tpu_inference/runner/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/compilation_manager.py +69 -35
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/kv_cache.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/kv_cache_manager.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/lora_utils.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/multimodal_manager.py +15 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/persistent_batch_manager.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/speculative_decoding_manager.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/structured_decoding_manager.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/tpu_runner.py +27 -9
- tpu_inference-0.13.2rc1/tpu_inference/spec_decode/__init__.py +13 -0
- tpu_inference-0.13.2rc1/tpu_inference/spec_decode/jax/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/spec_decode/jax/eagle3.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/tpu_info.py +14 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/utils.py +31 -30
- tpu_inference-0.13.2rc1/tpu_inference/worker/__init__.py +13 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/worker/tpu_worker.py +19 -6
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1/tpu_inference.egg-info}/PKG-INFO +1 -1
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference.egg-info/SOURCES.txt +90 -8
- tpu_inference-0.12.0.dev20251217/pyproject.toml +0 -0
- tpu_inference-0.12.0.dev20251217/tests/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tests/core/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tests/core/test_dp_scheduler.py +0 -899
- tpu_inference-0.12.0.dev20251217/tests/kernels/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tests/lora/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/core/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/core/sched/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/distributed/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/distributed/utils.py +0 -60
- tpu_inference-0.12.0.dev20251217/tpu_inference/executors/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/experimental/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/collectives/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/flash_attention/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/fused_moe/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/fused_moe/v1/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/megablox/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/mla/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/mla/v1/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/quantized_matmul/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v2/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py +0 -4147
- tpu_inference-0.12.0.dev20251217/tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py +0 -368
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/common/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/common/quant_methods.py +0 -8
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/attention/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/misc.py +0 -16
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/moe/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/jax/sample/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/lora/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/common/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/fp8_all_modules_w_only.yaml +0 -5
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/fp8_default.yaml +0 -6
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/int8_all_modules_w_only.yaml +0 -5
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/configs/int8_default.yaml +0 -6
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/jax/utils/quantization/mxfp4_utils.py +0 -105
- tpu_inference-0.12.0.dev20251217/tpu_inference/models/vllm/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/platforms/__init__.py +0 -2
- tpu_inference-0.12.0.dev20251217/tpu_inference/runner/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/spec_decode/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/spec_decode/jax/__init__.py +0 -0
- tpu_inference-0.12.0.dev20251217/tpu_inference/worker/__init__.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/LICENSE +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/README.md +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/requirements.txt +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/setup.cfg +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/core/test_core_tpu.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/core/test_disagg_executor.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/kernels/quantized_matmul_kernel_test.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/lora/utils.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/test_envs.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/core/core_tpu.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/core/disagg_executor.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/core/disagg_utils.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/env_override.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/envs.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/collectives/all_gather_matmul.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/collectives/util.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/flash_attention/kernel.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/quantized_matmul/kernel.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/quantized_matmul/util.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v2/kernel.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/layers/common/binary_search.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/logger.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/lora/torch_lora_ops.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/lora/torch_punica_tpu.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/block_table.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/input_batch.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference/runner/utils.py +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference.egg-info/dependency_links.txt +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference.egg-info/requires.txt +0 -0
- {tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tpu_inference.egg-info/top_level.txt +0 -0
|
@@ -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,25 @@
|
|
|
1
|
+
# This file contains additional dependencies needed for TPU v7x support.
|
|
2
|
+
# It is expected to be used in conjunction with the main requirements.txt file.
|
|
3
|
+
--pre
|
|
4
|
+
-i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
|
|
5
|
+
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
|
|
6
|
+
jax==0.8.1
|
|
7
|
+
jaxlib==0.8.1
|
|
8
|
+
jaxtyping==0.3.2
|
|
9
|
+
libtpu==0.0.31
|
|
10
|
+
|
|
11
|
+
tpu-info==0.7.1
|
|
12
|
+
yapf==0.43.0
|
|
13
|
+
pytest
|
|
14
|
+
pytest-mock
|
|
15
|
+
absl-py
|
|
16
|
+
numpy
|
|
17
|
+
google-cloud-storage
|
|
18
|
+
flax==0.11.1
|
|
19
|
+
torchax==0.0.10
|
|
20
|
+
qwix==0.1.1
|
|
21
|
+
torchvision==0.24.0
|
|
22
|
+
pathwaysutils
|
|
23
|
+
parameterized
|
|
24
|
+
numba==0.62.1
|
|
25
|
+
runai-model-streamer[s3,gcs]==0.15.0
|
|
@@ -20,26 +20,40 @@ def get_requirements() -> List[str]:
|
|
|
20
20
|
requirements = f.read().strip().split("\n")
|
|
21
21
|
resolved_requirements = []
|
|
22
22
|
for line in requirements:
|
|
23
|
+
if not line or line.startswith("#"):
|
|
24
|
+
continue
|
|
23
25
|
if line.startswith("-r "):
|
|
24
26
|
resolved_requirements += _read_requirements(line.split()[1])
|
|
25
|
-
elif line.startswith("--"):
|
|
27
|
+
elif line.startswith(("-", "--")):
|
|
26
28
|
continue
|
|
27
29
|
else:
|
|
28
30
|
resolved_requirements.append(line)
|
|
29
31
|
return resolved_requirements
|
|
30
32
|
|
|
31
33
|
try:
|
|
32
|
-
requirements = _read_requirements("
|
|
34
|
+
#requirements = _read_requirements("requirements_v7x.txt")
|
|
35
|
+
|
|
36
|
+
# For TPU v7x build
|
|
37
|
+
if os.getenv("IS_FOR_V7X", "false").lower() == "true":
|
|
38
|
+
print("Using requirements_v7x.txt")
|
|
39
|
+
requirements = _read_requirements("requirements_v7x.txt")
|
|
40
|
+
#requirements.extend(v7x_requirements)
|
|
41
|
+
else:
|
|
42
|
+
#For TPU v6e build
|
|
43
|
+
print("Using requirements.txt")
|
|
44
|
+
requirements = _read_requirements("requirements.txt")
|
|
45
|
+
|
|
33
46
|
except ValueError:
|
|
34
47
|
print("Failed to read requirements.txt in vllm_tpu.")
|
|
35
48
|
return requirements
|
|
36
49
|
|
|
37
50
|
|
|
38
51
|
def get_version():
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
52
|
+
version = os.getenv("VLLM_VERSION_OVERRIDE", "0.0.0").strip()
|
|
53
|
+
if os.getenv("IS_FOR_V7X", "false").lower() == "true":
|
|
54
|
+
version = f"{version}.post7"
|
|
42
55
|
|
|
56
|
+
return version
|
|
43
57
|
|
|
44
58
|
setup(
|
|
45
59
|
name="tpu_inference",
|
|
@@ -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.
|
{tpu_inference-0.12.0.dev20251217 → tpu_inference-0.13.2rc1}/tests/core/test_disagg_utils.py
RENAMED
|
@@ -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
|