LICENSE
MANIFEST.in
README.md
pyproject.toml
requirements.txt
setup.py
tests/__init__.py
tests/conftest.py
tests/test_base.py
tests/test_envs.py
tests/test_nightly_gate.py
tests/test_tpu_info.py
tests/test_utils.py
tests/test_video_input_plumbing.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_group_sampling_prefix_cache.py
tests/core/test_dp_scheduler.py
tests/core/test_init.py
tests/core/test_sched_utils.py
tests/distributed/__init__.py
tests/distributed/test_distributed_utils.py
tests/distributed/test_kv_transfer.py
tests/distributed/test_tpu_connector.py
tests/e2e/__init__.py
tests/e2e/mistral_small_4_correctness.py
tests/e2e/test_async_scheduler.py
tests/e2e/test_continue_decode.py
tests/e2e/test_data_parallel.py
tests/e2e/test_expert_parallel.py
tests/e2e/test_hybrid_kvcache.py
tests/e2e/test_model_loader.py
tests/e2e/test_multi_modal_inference.py
tests/e2e/test_pipeline_parallel.py
tests/e2e/test_rl_integration.py
tests/e2e/test_runai_model_streamer_loader.py
tests/e2e/test_sampling_params.py
tests/e2e/test_sequence_parallelism.py
tests/e2e/test_speculative_decoding.py
tests/e2e/test_step_pooling.py
tests/e2e/test_structured_decoding.py
tests/e2e/test_tensor_parallel.py
tests/executors/__init__.py
tests/executors/test_ray_distributed_executor.py
tests/executors/test_ray_distributed_executor_v2.py
tests/kernels/__init__.py
tests/kernels/causal_conv1d_test.py
tests/kernels/dense_gather_reduce_test.py
tests/kernels/flash_attention_kernel_test.py
tests/kernels/fused_gdn_kernel_test.py
tests/kernels/fused_moe_v1_test.py
tests/kernels/gather_reduce_test.py
tests/kernels/gdn_attention_v3_test.py
tests/kernels/gdn_decode_kernel_test.py
tests/kernels/gmm_test.py
tests/kernels/mla_tuned_vs_baseline_test.py
tests/kernels/mla_v2_test.py
tests/kernels/quantized_matmul_kernel_test.py
tests/kernels/ragged_gather_reduce_v2_test.py
tests/kernels/ragged_gather_v2_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/ragged_scatter_test.py
tests/kernels/recurrent_scan_test.py
tests/kernels/spmm_v1_test.py
tests/kernels/transpose_test.py
tests/kernels/collectives/__init__.py
tests/kernels/collectives/all_gather_matmul_kernel_test.py
tests/kernels/deepseek_v4/__init__.py
tests/kernels/deepseek_v4/compress_store_test.py
tests/kernels/deepseek_v4/csa_gather_test.py
tests/kernels/deepseek_v4/mla_swa_test.py
tests/kernels/deepseek_v4/mla_test.py
tests/kernels/deepseek_v4/rope_test.py
tests/kernels/deepseek_v4/sparse_mla_test.py
tests/kernels/deepseek_v4/test_streamindex_topk.py
tests/kernels/rpa_v3_cp/__init__.py
tests/kernels/rpa_v3_cp/ragged_paged_attention_kernel_cp_test.py
tests/kernels/stacked_rpa/__init__.py
tests/kernels/stacked_rpa/stacked_rpa_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_pcp_attention_interface.py
tests/layers/common/test_process_weights.py
tests/layers/common/test_quantization.py
tests/layers/common/test_ragged_conv1d.py
tests/layers/common/test_sharding.py
tests/layers/common/test_utils.py
tests/layers/common/utils.py
tests/layers/jax/__init__.py
tests/layers/jax/test_conv.py
tests/layers/jax/test_layers.py
tests/layers/jax/test_linear.py
tests/layers/jax/test_norm.py
tests/layers/jax/test_qwix.py
tests/layers/jax/test_rope.py
tests/layers/jax/test_rope_interface.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_gpt_oss_attention.py
tests/layers/jax/attention/test_llama4_attention.py
tests/layers/jax/moe/__init__.py
tests/layers/jax/moe/test_moe.py
tests/layers/jax/moe/test_moe_utils.py
tests/layers/jax/quantization/__init__.py
tests/layers/jax/quantization/test_compressed_tensors.py
tests/layers/jax/quantization/test_fp8.py
tests/layers/jax/quantization/test_mxfp4.py
tests/layers/jax/quantization/test_unquantized.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/nvfp4_utils.py
tests/layers/vllm/test_awq.py
tests/layers/vllm/test_compressed_tensors_moe.py
tests/layers/vllm/test_compressed_tensors_moe_w4a8.py
tests/layers/vllm/test_compressed_tensors_w4a4_nvfp4.py
tests/layers/vllm/test_compressed_tensors_w4a8_fp8.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_fused_moe.py
tests/layers/vllm/test_gdn_attention_op.py
tests/layers/vllm/test_mla_attention.py
tests/layers/vllm/test_moe.py
tests/layers/vllm/test_mxfp4.py
tests/layers/vllm/test_nvfp4.py
tests/layers/vllm/test_qkv_kv_replication.py
tests/layers/vllm/test_unquantized.py
tests/layers/vllm/ops/__init__.py
tests/layers/vllm/ops/test_scaled_dot_product_attention.py
tests/lora/__init__.py
tests/lora/conftest.py
tests/lora/test_bgmv.py
tests/lora/test_env_parsing.py
tests/lora/test_layers.py
tests/lora/test_lora.py
tests/lora/test_lora_adapter.py
tests/lora/test_lora_perf.py
tests/lora/test_punica_tpu.py
tests/lora/utils.py
tests/models/__init__.py
tests/models/common/__init__.py
tests/models/common/test_model_loader.py
tests/models/common/test_pathways_dummy_loader.py
tests/models/jax/__init__.py
tests/models/jax/conftest.py
tests/models/jax/test_deepseek_v3.py
tests/models/jax/test_dflash.py
tests/models/jax/test_gemma4.py
tests/models/jax/test_gemma4_kv_share.py
tests/models/jax/test_gemma4_mtp.py
tests/models/jax/test_gemma4_ple.py
tests/models/jax/test_gpt_oss.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_qwen3_dflash.py
tests/models/jax/test_qwen3_dflash_attention.py
tests/models/jax/test_weight_loading.py
tests/models/jax/utils/__init__.py
tests/models/jax/utils/test_multi_modal_utils.py
tests/models/jax/utils/test_weight_utils.py
tests/platforms/__init__.py
tests/platforms/test_tpu_platform.py
tests/runner/__init__.py
tests/runner/test_async_logprobs.py
tests/runner/test_block_table.py
tests/runner/test_continuous_block_pool.py
tests/runner/test_decode_loop.py
tests/runner/test_input_batch.py
tests/runner/test_kv_cache.py
tests/runner/test_kv_cache_manager.py
tests/runner/test_mm_encoder_jit_manager.py
tests/runner/test_moe_expert_ids.py
tests/runner/test_multimodal_manager.py
tests/runner/test_persistent_batch_manager.py
tests/runner/test_processed_logprobs.py
tests/runner/test_routed_experts_slots.py
tests/runner/test_speculative_decoding_manager.py
tests/runner/test_step_pooling_logic.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_dflash.py
tests/spec_decode/test_dflash_torchax.py
tests/spec_decode/test_eagle3.py
tests/spec_decode/test_utils.py
tests/worker/__init__.py
tests/worker/tpu_worker_test.py
tools/__init__.py
tools/kernel/__init__.py
tools/kernel/tuner/__init__.py
tools/kernel/tuner/v1/__init__.py
tools/kernel/tuner/v1/batched_rpa_kernel_tuner.py
tools/kernel/tuner/v1/example_kernel_tuner.py
tools/kernel/tuner/v1/flash_attention_kernel_tuner.py
tools/kernel/tuner/v1/inspect_result_cli.py
tools/kernel/tuner/v1/kernel_tuner_runner.py
tools/kernel/tuner/v1/mla_kernel_tuner.py
tools/kernel/tuner/v1/rpa_v3_kernel_tuner.py
tools/kernel/tuner/v1/utils.py
tools/kernel/tuner/v1/autotune/__init__.py
tools/kernel/tuner/v1/autotune/autotune_result_processing.py
tools/kernel/tuner/v1/autotune/bootstrap_kernel_tuners.py
tools/kernel/tuner/v1/autotune/kernel_autotune_config.py
tools/kernel/tuner/v1/autotune/update_tuned_params_template.py
tools/kernel/tuner/v1/common/__init__.py
tools/kernel/tuner/v1/common/kernel_tuner_base.py
tools/kernel/tuner/v1/common/tuner_datatypes.py
tools/kernel/tuner/v1/common/tuning_case_logger.py
tools/kernel/tuner/v1/common/utils.py
tools/kernel/tuner/v1/optimizer/__init__.py
tools/kernel/tuner/v1/optimizer/base_optimizer.py
tools/kernel/tuner/v1/optimizer/bayesian_optimizer.py
tools/kernel/tuner/v1/optimizer/sweep_optimizer.py
tools/kernel/tuner/v1/storage_management/__init__.py
tools/kernel/tuner/v1/storage_management/local_db_manager.py
tools/kernel/tuner/v1/storage_management/spanner_database_manager.py
tools/kernel/tuner/v1/storage_management/storage_manager.py
tools/kernel/tuner/v1/tests/__init__.py
tools/kernel/tuner/v1/tests/test_inspect_result_cli.py
tools/kernel/tuner/v1/tests/test_kernel_tuner_base.py
tools/kernel/tuner/v1/tests/test_kernel_tuner_runner.py
tools/kernel/tuner/v1/tests/test_optimizer.py
tools/kernel/tuner/v1/tests/test_tuned_params_structure.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/entry_points.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/core/sched/utils.py
tpu_inference/distributed/__init__.py
tpu_inference/distributed/host_kv_pool.py
tpu_inference/distributed/host_kv_pool_hma.py
tpu_inference/distributed/jax_parallel_state.py
tpu_inference/distributed/kv_transfer.py
tpu_inference/distributed/tpu_connector.py
tpu_inference/distributed/tpu_connector_hma.py
tpu_inference/distributed/tpu_connector_stats.py
tpu_inference/distributed/tpu_raiden_connector.py
tpu_inference/distributed/transfer_stats.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/executors/ray_distributed_executor_v2.py
tpu_inference/kernels/__init__.py
tpu_inference/kernels/causal_conv1d/__init__.py
tpu_inference/kernels/causal_conv1d/causal_conv1d.py
tpu_inference/kernels/causal_conv1d/strided_ldst.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/collectives/hierrs_sc/__init__.py
tpu_inference/kernels/collectives/hierrs_sc/config.py
tpu_inference/kernels/collectives/hierrs_sc/dma_pipeline.py
tpu_inference/kernels/collectives/hierrs_sc/kernel.py
tpu_inference/kernels/collectives/hierrs_sc/topology.py
tpu_inference/kernels/collectives/hierrs_sc/wrapper.py
tpu_inference/kernels/experimental/__init__.py
tpu_inference/kernels/experimental/batched_rpa/__init__.py
tpu_inference/kernels/experimental/batched_rpa/bref_override.py
tpu_inference/kernels/experimental/batched_rpa/configs.py
tpu_inference/kernels/experimental/batched_rpa/flash_attention.py
tpu_inference/kernels/experimental/batched_rpa/kernel.py
tpu_inference/kernels/experimental/batched_rpa/schedule.py
tpu_inference/kernels/experimental/batched_rpa/stitch_utils.py
tpu_inference/kernels/experimental/batched_rpa/tuned_params.py
tpu_inference/kernels/experimental/batched_rpa/utils.py
tpu_inference/kernels/experimental/batched_rpa/wrapper.py
tpu_inference/kernels/experimental/deepseek_v4/__init__.py
tpu_inference/kernels/experimental/deepseek_v4/rope.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/__init__.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/buffered_ref.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/compress_store_ref.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/compressor_v1.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/compute.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/config.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/kernel.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/proj_and_save_state.py
tpu_inference/kernels/experimental/deepseek_v4/compress_and_store/project_and_save_state_ref.py
tpu_inference/kernels/experimental/deepseek_v4/core_attention/__init__.py
tpu_inference/kernels/experimental/deepseek_v4/core_attention/csa_gather.py
tpu_inference/kernels/experimental/deepseek_v4/core_attention/mla.py
tpu_inference/kernels/experimental/deepseek_v4/core_attention/mla_swa.py
tpu_inference/kernels/experimental/deepseek_v4/core_attention/sparse_mla.py
tpu_inference/kernels/experimental/deepseek_v4/indexer/__init__.py
tpu_inference/kernels/experimental/deepseek_v4/indexer/streamindex_topk.py
tpu_inference/kernels/experimental/rpa_v3_cp/__init__.py
tpu_inference/kernels/experimental/rpa_v3_cp/kernel.py
tpu_inference/kernels/experimental/stacked_rpa/__init__.py
tpu_inference/kernels/experimental/stacked_rpa/bref_override.py
tpu_inference/kernels/experimental/stacked_rpa/configs.py
tpu_inference/kernels/experimental/stacked_rpa/envvars.py
tpu_inference/kernels/experimental/stacked_rpa/flash_attention.py
tpu_inference/kernels/experimental/stacked_rpa/flash_attention_chunked.py
tpu_inference/kernels/experimental/stacked_rpa/kernel.py
tpu_inference/kernels/experimental/stacked_rpa/schedule.py
tpu_inference/kernels/experimental/stacked_rpa/stitch_utils.py
tpu_inference/kernels/experimental/stacked_rpa/tuned_block_sizes.py
tpu_inference/kernels/experimental/stacked_rpa/utils.py
tpu_inference/kernels/experimental/stacked_rpa/wrapper.py
tpu_inference/kernels/flash_attention/__init__.py
tpu_inference/kernels/flash_attention/kernel.py
tpu_inference/kernels/flash_attention/tuned_params.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/gdn/__init__.py
tpu_inference/kernels/gdn/triangle_solver.py
tpu_inference/kernels/gdn/reference/__init__.py
tpu_inference/kernels/gdn/reference/ragged_gated_delta_rule_chunked.py
tpu_inference/kernels/gdn/reference/ragged_gated_delta_rule_ref.py
tpu_inference/kernels/gdn/reference/ragged_gated_delta_rule_wrapper.py
tpu_inference/kernels/gdn/v1/__init__.py
tpu_inference/kernels/gdn/v1/fused_gdn_decode_kernel.py
tpu_inference/kernels/gdn/v1/fused_gdn_kernel_common.py
tpu_inference/kernels/gdn/v1/fused_gdn_kernel_wrapper.py
tpu_inference/kernels/gdn/v1/fused_gdn_recurrent_kernel.py
tpu_inference/kernels/gdn/v2/__init__.py
tpu_inference/kernels/gdn/v2/compute_schedule_v2.py
tpu_inference/kernels/gdn/v2/gdn_decode_kernel.py
tpu_inference/kernels/gdn/v2/recurrent_scan_impl.py
tpu_inference/kernels/gdn/v2/recurrent_scan_v2.py
tpu_inference/kernels/gdn/v3/__init__.py
tpu_inference/kernels/gdn/v3/compute_conv1d.py
tpu_inference/kernels/gdn/v3/compute_gdn.py
tpu_inference/kernels/gdn/v3/config.py
tpu_inference/kernels/gdn/v3/memory_ref.py
tpu_inference/kernels/gdn/v3/metadata.py
tpu_inference/kernels/gdn/v3/vmem_ldst.py
tpu_inference/kernels/gdn/v3/wrapper.py
tpu_inference/kernels/megablox/__init__.py
tpu_inference/kernels/megablox/common.py
tpu_inference/kernels/megablox/gmm.py
tpu_inference/kernels/megablox/gmm_v2.py
tpu_inference/kernels/megablox/tuned_block_sizes.py
tpu_inference/kernels/mla/__init__.py
tpu_inference/kernels/mla/v1/__init__.py
tpu_inference/kernels/mla/v1/kernel.py
tpu_inference/kernels/mla/v2/__init__.py
tpu_inference/kernels/mla/v2/kernel.py
tpu_inference/kernels/mla/v2/kv_utils.py
tpu_inference/kernels/mla/v2/transpose.py
tpu_inference/kernels/mla/v2/tuned_params.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/kernels/sparse_core/__init__.py
tpu_inference/kernels/sparse_core/core_map_helper.py
tpu_inference/kernels/sparse_core/dense_gather_reduce.py
tpu_inference/kernels/sparse_core/gather_reduce.py
tpu_inference/kernels/sparse_core/ragged_gather_reduce_v2.py
tpu_inference/kernels/sparse_core/ragged_gather_v2.py
tpu_inference/kernels/sparse_core/ragged_scatter.py
tpu_inference/kernels/structured_sparse_matmul/__init__.py
tpu_inference/kernels/structured_sparse_matmul/v1/__init__.py
tpu_inference/kernels/structured_sparse_matmul/v1/spmm.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/cp_attention.py
tpu_inference/layers/common/dflash_attention_interface.py
tpu_inference/layers/common/fused_moe_gmm.py
tpu_inference/layers/common/gdn_attention.py
tpu_inference/layers/common/linear.py
tpu_inference/layers/common/moe.py
tpu_inference/layers/common/quant_methods.py
tpu_inference/layers/common/ragged_conv1d_jax.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/process_weights/moe_weights.py
tpu_inference/layers/common/quantization/__init__.py
tpu_inference/layers/common/quantization/configs.py
tpu_inference/layers/common/quantization/fp8.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/conv.py
tpu_inference/layers/jax/embed.py
tpu_inference/layers/jax/layers.py
tpu_inference/layers/jax/linear.py
tpu_inference/layers/jax/misc.py
tpu_inference/layers/jax/norm.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/gpt_oss_attention.py
tpu_inference/layers/jax/attention/llama4_attention.py
tpu_inference/layers/jax/moe/__init__.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/compressed_tensors.py
tpu_inference/layers/jax/quantization/configs.py
tpu_inference/layers/jax/quantization/fp8.py
tpu_inference/layers/jax/quantization/mxfp4.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/backends/__init__.py
tpu_inference/layers/vllm/backends/flash_attn.py
tpu_inference/layers/vllm/backends/flash_attn_mla.py
tpu_inference/layers/vllm/custom_ops/__init__.py
tpu_inference/layers/vllm/custom_ops/embedding.py
tpu_inference/layers/vllm/custom_ops/fused_moe.py
tpu_inference/layers/vllm/custom_ops/gdn_attention_op.py
tpu_inference/layers/vllm/custom_ops/linear.py
tpu_inference/layers/vllm/custom_ops/mhc.py
tpu_inference/layers/vllm/custom_ops/mla_attention.py
tpu_inference/layers/vllm/custom_ops/rope.py
tpu_inference/layers/vllm/custom_ops/experimental/__init__.py
tpu_inference/layers/vllm/custom_ops/experimental/deepseek_v4/__init__.py
tpu_inference/layers/vllm/custom_ops/experimental/deepseek_v4/deepseek_v4_attention.py
tpu_inference/layers/vllm/custom_ops/experimental/deepseek_v4/deepseek_v4_compressor.py
tpu_inference/layers/vllm/custom_ops/experimental/deepseek_v4/deepseek_v4_indexer.py
tpu_inference/layers/vllm/interface/__init__.py
tpu_inference/layers/vllm/interface/moe.py
tpu_inference/layers/vllm/ops/__init__.py
tpu_inference/layers/vllm/ops/scaled_dot_product_attention.py
tpu_inference/layers/vllm/process_weights/__init__.py
tpu_inference/layers/vllm/process_weights/cleanup_sharding.py
tpu_inference/layers/vllm/quantization/__init__.py
tpu_inference/layers/vllm/quantization/awq.py
tpu_inference/layers/vllm/quantization/base.py
tpu_inference/layers/vllm/quantization/configs.py
tpu_inference/layers/vllm/quantization/deepseek_v4_fp8.py
tpu_inference/layers/vllm/quantization/fp8.py
tpu_inference/layers/vllm/quantization/mxfp4.py
tpu_inference/layers/vllm/quantization/nvfp4.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/__init__.py
tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8.py
tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.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/lora_manager.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/compiler_options.py
tpu_inference/models/common/interface.py
tpu_inference/models/common/kv_share.py
tpu_inference/models/common/model_loader.py
tpu_inference/models/common/pathways_dummy_loader.py
tpu_inference/models/jax/__init__.py
tpu_inference/models/jax/deepseek_v3.py
tpu_inference/models/jax/dflash.py
tpu_inference/models/jax/gemma4.py
tpu_inference/models/jax/gemma4_mm.py
tpu_inference/models/jax/gemma4_mtp.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/qwen3_dflash.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/dflash.py
tpu_inference/models/vllm/vllm_model_loader.py
tpu_inference/models/vllm/vllm_model_wrapper.py
tpu_inference/models/vllm/vllm_model_wrapper_context.py
tpu_inference/models/vllm/experimental/__init__.py
tpu_inference/models/vllm/experimental/deepseek_v4.py
tpu_inference/models/vllm/experimental/model_patcher.py
tpu_inference/models/vllm/experimental/qwen3_omni_patcher.py
tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py
tpu_inference/models/vllm/experimental/vision_tower_jit.py
tpu_inference/offload/__init__.py
tpu_inference/offload/cpu_backend.py
tpu_inference/offload/metrics.py
tpu_inference/offload/offload_manager.py
tpu_inference/offload/raiden_offload_connector.py
tpu_inference/offload/tpu_offload_connector.py
tpu_inference/offload/utils.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/continuous_block_pool.py
tpu_inference/runner/decode_loop.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/mm_encoder_jit_manager.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/dflash.py
tpu_inference/spec_decode/jax/eagle3.py
tpu_inference/spec_decode/jax/utils.py
tpu_inference/spec_decode/vllm/__init__.py
tpu_inference/spec_decode/vllm/dflash.py
tpu_inference/worker/__init__.py
tpu_inference/worker/tpu_worker.py