LICENSE
MANIFEST.in
README.md
pyproject.toml
requirements.txt
setup.py
tests/__init__.py
tests/test_base.py
tests/test_envs.py
tests/test_quantization.py
tests/test_tpu_info.py
tests/test_utils.py
tests/core/__init__.py
tests/core/test_core_tpu.py
tests/core/test_disagg_executor.py
tests/core/test_disagg_utils.py
tests/core/test_dp_scheduler.py
tests/core/test_init.py
tests/kernels/__init__.py
tests/kernels/fused_moe_v1_test.py
tests/kernels/gmm_test.py
tests/kernels/mla_v1_test.py
tests/kernels/quantized_matmul_kernel_test.py
tests/kernels/ragged_kv_cache_update_v2_test.py
tests/kernels/ragged_paged_attention_kernel_v2_test.py
tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py
tests/kernels/ragged_paged_attention_kernel_v3_test.py
tests/lora/__init__.py
tests/lora/conftest.py
tests/lora/test_bgmv.py
tests/lora/test_layers.py
tests/lora/test_lora.py
tests/lora/test_lora_perf.py
tests/lora/utils.py
tpu_inference/__init__.py
tpu_inference/env_override.py
tpu_inference/envs.py
tpu_inference/logger.py
tpu_inference/tpu_info.py
tpu_inference/utils.py
tpu_inference.egg-info/PKG-INFO
tpu_inference.egg-info/SOURCES.txt
tpu_inference.egg-info/dependency_links.txt
tpu_inference.egg-info/requires.txt
tpu_inference.egg-info/top_level.txt
tpu_inference/core/__init__.py
tpu_inference/core/core_tpu.py
tpu_inference/core/disagg_executor.py
tpu_inference/core/disagg_utils.py
tpu_inference/core/sched/__init__.py
tpu_inference/core/sched/dp_scheduler.py
tpu_inference/distributed/__init__.py
tpu_inference/distributed/jax_parallel_state.py
tpu_inference/distributed/tpu_connector.py
tpu_inference/distributed/utils.py
tpu_inference/executors/__init__.py
tpu_inference/executors/ray_distributed_executor.py
tpu_inference/experimental/__init__.py
tpu_inference/experimental/llama3_jax_stashed.py
tpu_inference/kernels/__init__.py
tpu_inference/kernels/collectives/__init__.py
tpu_inference/kernels/collectives/all_gather_matmul.py
tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py
tpu_inference/kernels/collectives/util.py
tpu_inference/kernels/flash_attention/__init__.py
tpu_inference/kernels/flash_attention/kernel.py
tpu_inference/kernels/fused_moe/__init__.py
tpu_inference/kernels/fused_moe/v1/__init__.py
tpu_inference/kernels/fused_moe/v1/kernel.py
tpu_inference/kernels/megablox/__init__.py
tpu_inference/kernels/megablox/common.py
tpu_inference/kernels/megablox/gmm.py
tpu_inference/kernels/mla/__init__.py
tpu_inference/kernels/mla/v1/__init__.py
tpu_inference/kernels/mla/v1/kernel.py
tpu_inference/kernels/quantized_matmul/__init__.py
tpu_inference/kernels/quantized_matmul/kernel.py
tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py
tpu_inference/kernels/quantized_matmul/util.py
tpu_inference/kernels/ragged_paged_attention/__init__.py
tpu_inference/kernels/ragged_paged_attention/v2/__init__.py
tpu_inference/kernels/ragged_paged_attention/v2/kernel.py
tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py
tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py
tpu_inference/kernels/ragged_paged_attention/v3/__init__.py
tpu_inference/kernels/ragged_paged_attention/v3/kernel.py
tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py
tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py
tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py
tpu_inference/kernels/ragged_paged_attention/v3/util.py
tpu_inference/layers/__init__.py
tpu_inference/layers/common/__init__.py
tpu_inference/layers/common/attention_interface.py
tpu_inference/layers/common/attention_metadata.py
tpu_inference/layers/common/binary_search.py
tpu_inference/layers/common/quant_methods.py
tpu_inference/layers/common/sharding.py
tpu_inference/layers/jax/__init__.py
tpu_inference/layers/jax/base.py
tpu_inference/layers/jax/constants.py
tpu_inference/layers/jax/layers.py
tpu_inference/layers/jax/misc.py
tpu_inference/layers/jax/pp_utils.py
tpu_inference/layers/jax/rope.py
tpu_inference/layers/jax/rope_interface.py
tpu_inference/layers/jax/transformer_block.py
tpu_inference/layers/jax/attention/__init__.py
tpu_inference/layers/jax/attention/attention.py
tpu_inference/layers/jax/attention/deepseek_v3_attention.py
tpu_inference/layers/jax/attention/gpt_oss_attention.py
tpu_inference/layers/jax/attention/llama4_attention.py
tpu_inference/layers/jax/moe/__init__.py
tpu_inference/layers/jax/moe/deepseek_v3_moe.py
tpu_inference/layers/jax/moe/gpt_oss_moe.py
tpu_inference/layers/jax/moe/moe.py
tpu_inference/layers/jax/sample/__init__.py
tpu_inference/layers/jax/sample/rejection_sampler.py
tpu_inference/layers/jax/sample/sampling.py
tpu_inference/layers/jax/sample/sampling_metadata.py
tpu_inference/layers/vllm/__init__.py
tpu_inference/layers/vllm/attention.py
tpu_inference/layers/vllm/fused_moe.py
tpu_inference/layers/vllm/linear_common.py
tpu_inference/layers/vllm/sharding.py
tpu_inference/layers/vllm/quantization/__init__.py
tpu_inference/layers/vllm/quantization/awq.py
tpu_inference/layers/vllm/quantization/common.py
tpu_inference/layers/vllm/quantization/mxfp4.py
tpu_inference/layers/vllm/quantization/unquantized.py
tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py
tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py
tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
tpu_inference/lora/__init__.py
tpu_inference/lora/torch_lora_ops.py
tpu_inference/lora/torch_punica_tpu.py
tpu_inference/models/__init__.py
tpu_inference/models/common/__init__.py
tpu_inference/models/common/model_loader.py
tpu_inference/models/jax/__init__.py
tpu_inference/models/jax/deepseek_v3.py
tpu_inference/models/jax/gpt_oss.py
tpu_inference/models/jax/jax_intermediate_tensor.py
tpu_inference/models/jax/llama3.py
tpu_inference/models/jax/llama4.py
tpu_inference/models/jax/llama_eagle3.py
tpu_inference/models/jax/llama_guard_4.py
tpu_inference/models/jax/qwen2.py
tpu_inference/models/jax/qwen2_5_vl.py
tpu_inference/models/jax/qwen3.py
tpu_inference/models/jax/utils/__init__.py
tpu_inference/models/jax/utils/file_utils.py
tpu_inference/models/jax/utils/multi_modal_utils.py
tpu_inference/models/jax/utils/weight_utils.py
tpu_inference/models/jax/utils/quantization/__init__.py
tpu_inference/models/jax/utils/quantization/mxfp4_utils.py
tpu_inference/models/jax/utils/quantization/quantization_utils.py
tpu_inference/models/jax/utils/quantization/configs/fp8_all_modules_w_only.yaml
tpu_inference/models/jax/utils/quantization/configs/fp8_default.yaml
tpu_inference/models/jax/utils/quantization/configs/int8_all_modules_w_only.yaml
tpu_inference/models/jax/utils/quantization/configs/int8_default.yaml
tpu_inference/models/vllm/__init__.py
tpu_inference/models/vllm/vllm_model_wrapper.py
tpu_inference/models/vllm/vllm_model_wrapper_context.py
tpu_inference/platforms/__init__.py
tpu_inference/platforms/tpu_platform.py
tpu_inference/runner/__init__.py
tpu_inference/runner/block_table.py
tpu_inference/runner/compilation_manager.py
tpu_inference/runner/input_batch.py
tpu_inference/runner/kv_cache.py
tpu_inference/runner/kv_cache_manager.py
tpu_inference/runner/lora_utils.py
tpu_inference/runner/multimodal_manager.py
tpu_inference/runner/persistent_batch_manager.py
tpu_inference/runner/speculative_decoding_manager.py
tpu_inference/runner/structured_decoding_manager.py
tpu_inference/runner/tpu_runner.py
tpu_inference/runner/utils.py
tpu_inference/spec_decode/__init__.py
tpu_inference/spec_decode/jax/__init__.py
tpu_inference/spec_decode/jax/eagle3.py
tpu_inference/worker/__init__.py
tpu_inference/worker/tpu_worker.py