LICENSE
MANIFEST.in
README.md
pyproject.toml
requirements.txt
setup.py
tests/__init__.py
tests/test_base.py
tests/test_envs.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/distributed/__init__.py
tests/distributed/test_distributed_utils.py
tests/distributed/test_tpu_connector.py
tests/e2e/__init__.py
tests/e2e/test_async_scheduler.py
tests/e2e/test_data_parallel.py
tests/e2e/test_hybrid_kvcache.py
tests/e2e/test_local_disagg.py
tests/e2e/test_model_loader.py
tests/e2e/test_multi_modal_inference.py
tests/e2e/test_pipeline_parallel.py
tests/e2e/test_runai_model_streamer_loader.py
tests/e2e/test_sampling_params.py
tests/e2e/test_speculative_decoding.py
tests/e2e/test_structured_decoding.py
tests/executors/__init__.py
tests/executors/test_ray_distributed_executor.py
tests/experimental/__init__.py
tests/experimental/test_llama3_jax_stashed.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/kernels/collectives/__init__.py
tests/kernels/collectives/all_gather_matmul_kernel_test.py
tests/layers/__init__.py
tests/layers/test_jax.py
tests/layers/common/__init__.py
tests/layers/common/test_attention_interface.py
tests/layers/common/test_quantization.py
tests/layers/jax/__init__.py
tests/layers/jax/test_layers.py
tests/layers/jax/test_linear.py
tests/layers/jax/test_qwix.py
tests/layers/jax/test_rope.py
tests/layers/jax/test_sharding.py
tests/layers/jax/test_transformer_block.py
tests/layers/jax/attention/__init__.py
tests/layers/jax/attention/test_common_attention.py
tests/layers/jax/attention/test_deepseek_v3_attention.py
tests/layers/jax/attention/test_llama4_attention.py
tests/layers/jax/moe/__init__.py
tests/layers/jax/moe/test_deepseek_moe.py
tests/layers/jax/sample/__init__.py
tests/layers/jax/sample/test_rejection_sampler.py
tests/layers/jax/sample/test_sampling.py
tests/layers/jax/sample/test_sampling_metadata.py
tests/layers/vllm/__init__.py
tests/layers/vllm/test_attention.py
tests/layers/vllm/test_awq.py
tests/layers/vllm/test_compressed_tensors_moe.py
tests/layers/vllm/test_compressed_tensors_w8a8_fp8.py
tests/layers/vllm/test_compressed_tensors_w8a8_int8.py
tests/layers/vllm/test_fp8.py
tests/layers/vllm/test_mxfp4.py
tests/layers/vllm/test_unquantized.py
tests/layers/vllm/utils.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
tests/models/__init__.py
tests/models/common/__init__.py
tests/models/common/test_model_loader.py
tests/models/jax/__init__.py
tests/models/jax/test_deepseek_v3.py
tests/models/jax/test_llama3.py
tests/models/jax/test_llama4.py
tests/models/jax/test_llama_eagle3.py
tests/models/jax/test_llama_guard_4.py
tests/models/jax/test_qwen2.py
tests/models/jax/test_qwen2_5_vl.py
tests/models/jax/test_qwen3.py
tests/models/jax/test_weight_loading.py
tests/models/jax/utils/__init__.py
tests/models/jax/utils/test_multi_modal_utils.py
tests/platforms/__init__.py
tests/platforms/test_tpu_platform.py
tests/runner/__init__.py
tests/runner/test_block_table.py
tests/runner/test_input_batch.py
tests/runner/test_kv_cache.py
tests/runner/test_kv_cache_manager.py
tests/runner/test_multimodal_manager.py
tests/runner/test_persistent_batch_manager.py
tests/runner/test_speculative_decoding_manager.py
tests/runner/test_structured_decoding_manager.py
tests/runner/test_tpu_runner.py
tests/runner/test_tpu_runner_dp.py
tests/runner/test_tpu_runner_mesh.py
tests/runner/test_utils.py
tests/spec_decode/__init__.py
tests/spec_decode/test_eagle3.py
tests/worker/__init__.py
tests/worker/tpu_worker_test.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/multiproc_executor.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/fused_moe/v1/tuned_block_sizes.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/blockwise_kernel.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/fused_moe_gmm.py
tpu_inference/layers/common/quant_methods.py
tpu_inference/layers/common/sharding.py
tpu_inference/layers/common/utils.py
tpu_inference/layers/common/process_weights/__init__.py
tpu_inference/layers/common/process_weights/linear_weights.py
tpu_inference/layers/common/quantization/__init__.py
tpu_inference/layers/common/quantization/configs.py
tpu_inference/layers/common/quantization/unquantized.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/linear.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/dense_moe.py
tpu_inference/layers/jax/moe/gpt_oss_moe.py
tpu_inference/layers/jax/moe/moe.py
tpu_inference/layers/jax/moe/sparse_moe.py
tpu_inference/layers/jax/moe/utils.py
tpu_inference/layers/jax/quantization/__init__.py
tpu_inference/layers/jax/quantization/configs.py
tpu_inference/layers/jax/quantization/unquantized.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.py
tpu_inference/layers/vllm/process_weights/__init__.py
tpu_inference/layers/vllm/process_weights/cleanup_sharding.py
tpu_inference/layers/vllm/process_weights/fused_moe_weights.py
tpu_inference/layers/vllm/quantization/__init__.py
tpu_inference/layers/vllm/quantization/awq.py
tpu_inference/layers/vllm/quantization/configs.py
tpu_inference/layers/vllm/quantization/fp8.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/qwix/__init__.py
tpu_inference/models/jax/utils/qwix/qwix_utils.py
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