tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511130813__py3-none-any.whl

{tpu_inference-0.0.1rc1.dist-info → tpu_inference-0.11.1.dev202511130813.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tpu_inference
-Version: 0.0.1rc1
+Version: 0.11.1.dev202511130813
 Author: tpu_inference Contributors
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
@@ -14,7 +14,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: tpu-info==0.7.1
+Requires-Dist: tpu-info==0.4.0
 Requires-Dist: yapf==0.43.0
 Requires-Dist: pytest
 Requires-Dist: pytest-mock
@@ -27,11 +27,10 @@ Requires-Dist: jaxtyping
 Requires-Dist: flax==0.11.1
 Requires-Dist: torchax==0.0.7
 Requires-Dist: qwix==0.1.1
-Requires-Dist: torchvision==0.24.0
+Requires-Dist: torchvision==0.23.0
 Requires-Dist: pathwaysutils
 Requires-Dist: parameterized
 Requires-Dist: numba==0.62.1
-Requires-Dist: runai-model-streamer[gcs,s3]==0.15.0
 Dynamic: author
 Dynamic: classifier
 Dynamic: description

{tpu_inference-0.0.1rc1.dist-info → tpu_inference-0.11.1.dev202511130813.dist-info}/RECORD

@@ -1,9 +1,8 @@
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/test_base.py,sha256=Ct5WFRMHL7IHEIxk8FrzAvO8m0xFuDpzDBKkAKKAL2Q,7341
-tests/test_envs.py,sha256=h502VxL2gvhECm8u5uDh5JTGvhFf_DfQO88SpqOFMzE,7135
 tests/test_quantization.py,sha256=IT5ASyS1uuWcxc22kRtBcA-V4j3Z3hb7pMztm3GOlBs,34445
 tests/test_tpu_info.py,sha256=ZrwlMsp8ffITkS_b8Q1t_QG-a-WVAd4NUcjHhGibcsI,4670
-tests/test_utils.py,sha256=GIXLdd-x4gnqSLrySXGk22phqPc8MegFd7ph1Jj8OcU,8182
+tests/test_utils.py,sha256=szRg4UB36RcgIvbEd9xMhKYbWi-O4XAUWGJlIU6FJ9E,7983
 tests/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/core/test_core_tpu.py,sha256=r496rk1eOsK_F4nvm9zprl_T-RcO6eCUb7LuVReOZno,21413
 tests/core/test_disagg_executor.py,sha256=QdE2YZs08EyDDCmSjhiXkXqQ9BJTgO6csr_E1xkkfSg,2256
@@ -11,37 +10,37 @@ tests/core/test_disagg_utils.py,sha256=alktTGppaGdg-_un0Amz8Y0IDQz-xNJN0dXG-YApE
 tests/core/test_dp_scheduler.py,sha256=IwCR1Vs96V4CQdWA051rNaYxxr2V_byA1yx9HWyRoMg,37339
 tests/core/test_init.py,sha256=NEFI5A9eKGu4rmeJ2iqd0EmhlA3bzbVkXmMi1PV1b9U,1687
 tests/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/kernels/fused_moe_v1_test.py,sha256=sQ6gvpti94fpPYrSZn7frPPNjqbVmRibFtenVrGGCA4,10403
+tests/kernels/fused_moe_v1_test.py,sha256=c6zbSHQDzOseeyL9VCjQeP7zayNnwYf059CPlKcvZzQ,3137
 tests/kernels/mla_v1_test.py,sha256=oZc4TCgquiG0KOeWfv46yJbUIpro_CgCMFc7vzyB7t8,11646
 tests/kernels/quantized_matmul_kernel_test.py,sha256=od5-zXFjcsc_gWGRDrREL8E_ftymNniQVTzgtkBo_Gc,5679
 tests/kernels/ragged_kv_cache_update_v2_test.py,sha256=6-HjP5CoUG-kcuP8MS-JJVMiBnPRo_zadS3VInnO0D4,10821
 tests/kernels/ragged_paged_attention_kernel_v2_test.py,sha256=pWqo9UYF0tzwgBKO_xYw-TYSPrtAsKcMK5Haj8hFG7I,11340
-tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py,sha256=vLfe1I_vLdf0SqtBuBL7QHLSklrhWOOzYF-I_I3rdNo,16309
+tests/kernels/ragged_paged_attention_kernel_v3_hd64_test.py,sha256=JhIElqUZIRqIsfQ3U1RUzSiH_gz_SabAqDosGGZ2tlA,16321
 tests/kernels/ragged_paged_attention_kernel_v3_test.py,sha256=Hrd8iUkS1pS3rxeTyY53aYRg_ZL_d3NqgBXvOgnigSU,14838
 tests/lora/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/lora/conftest.py,sha256=EXjwE1CjmUUlMEXpyE3UwxvgrKUllE73I8BNKfP1FTc,984
 tests/lora/test_bgmv.py,sha256=gQxWsJdNX2nkrE2xyrG0exwf3E2eHm2k2nkEXoANuQc,1359
-tests/lora/test_layers.py,sha256=6B4HhMAItQmt0hPAQgyXgwSYs7b3bIbUf6LaPsqXLzY,25923
+tests/lora/test_layers.py,sha256=21ekYlsK36r1GPZOfzs7E-KIsfI1JcuZl1E6vaQbHf4,26273
 tests/lora/test_lora.py,sha256=wJiF1P1BDnPN8TLX2tlFtdZ_QCkV-S9nPl6_uR6DqFc,4439
-tests/lora/utils.py,sha256=rY0tDZEZe58ye4-ykwrTnsiWuLcaEG57N_Rua90bDXI,2726
-tpu_inference/__init__.py,sha256=p4MaepRdN7723FUNE-3pOMxZWjFn4_TVFgjrNyty4JE,2304
+tests/lora/utils.py,sha256=dR_v1H20vPVjFHdBhDajWOz0WJZlKuPLgMFQsME0LtA,3009
+tpu_inference/__init__.py,sha256=7IduGWw-_fwx0VA6EvC_AqHF67fnnShz6YvkqCfvFx8,1317
 tpu_inference/env_override.py,sha256=pmL7lfs_rGCP92ya3wuWuudsCYeOMZ6tFZY82A4KkQc,365
-tpu_inference/envs.py,sha256=ugze6VdQ_hG1IxUCbcgXZq7a22fZ-Lora3V_fkFOefw,5714
+tpu_inference/envs.py,sha256=MTT_Pdtd6cAcciYjv1OekEmvspaq3SYL0oR_jDkQ_aE,3948
 tpu_inference/logger.py,sha256=HQCz7NefmbturuhOC7-3Ixbtcdgoz4g9FHh2RB6o8cc,334
-tpu_inference/tpu_info.py,sha256=3iilHRQSFjwMJwhKcuuawTm7mhwkgHbj4zi6CiAySrs,2265
-tpu_inference/utils.py,sha256=mHbjI8fxInPxagLsSUg-R3DzSz-X7WYNdoorPYoE3hg,10855
+tpu_inference/tpu_info.py,sha256=9UohshkndR6dZpGWpWXfTD4qvIVdVgHf0yOoSEkLTrw,2276
+tpu_inference/utils.py,sha256=LWEshJgUdB20H2fDA-QI-Sk4EP7PD_FWvW3Mrqb-k8M,10054
 tpu_inference/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/core/core_tpu.py,sha256=WDD3koE_j1QhWS2BbMA2aQOZayPZm4tYPvzL4YCX2jY,33294
+tpu_inference/core/core_tpu.py,sha256=JdN4-xaxSWnzY4T181SCnbZ5HEnwQ5IifYA9ybF4pWo,32710
 tpu_inference/core/disagg_executor.py,sha256=HZpgYMVxRxm0RQxO4l8IDYBWJ6Z3Tac6xavc5otcirc,4657
-tpu_inference/core/disagg_utils.py,sha256=lv8MAVoAjtcmTaenUXVokg2q3d0tzsma86UiQlQ3omY,1492
+tpu_inference/core/disagg_utils.py,sha256=ufWNFWQ5n4YnZpPOtoReHlYo4dlN7AbIqCyqS4an0t4,1572
 tpu_inference/core/sched/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/core/sched/dp_scheduler.py,sha256=mKs8Ms46szdlBfo8hjdqis2ZKAZbcKnHAGfEr0X5R8g,22527
 tpu_inference/distributed/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/distributed/jax_parallel_state.py,sha256=5_xCwcL03lFPUoSO_OP7hIVKpUFroW1m-jVO7R6FbUc,2223
-tpu_inference/distributed/tpu_connector.py,sha256=kLaTwy6BrAThJeFkd1soJ47bBo5iGp4GjUJs7xFx4Tg,29696
-tpu_inference/distributed/utils.py,sha256=1KIREn28Zg10O-MSUkVQMRzS09WoGc_VLGOX4QTFJac,1504
+tpu_inference/distributed/tpu_connector.py,sha256=Zah46Sm5iOuh72SzXw69NxMc0MLnqsLEpe2BfDhpnqA,29731
+tpu_inference/distributed/utils.py,sha256=RwFQi8G4TzN1g9RjQu0pb5JxSc_jhoIZVsFJo0uHjxo,1513
 tpu_inference/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/executors/ray_distributed_executor.py,sha256=9CnzWb8aurH1B0tJfMHB73F-RQBGqSf5DnymetBvZ5o,16225
+tpu_inference/executors/ray_distributed_executor.py,sha256=UgJP-XSgDPKDj_mkVQ16XrRN96juVpnFl6fdWEyFL_Y,15249
 tpu_inference/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/experimental/llama3_jax_stashed.py,sha256=YK1oSIfto9ALo-HB45XfSrbq9XgVbE4m2C-9zRwmSzI,10913
 tpu_inference/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -53,7 +52,7 @@ tpu_inference/kernels/flash_attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCe
 tpu_inference/kernels/flash_attention/kernel.py,sha256=n8gmAFVfchMXlyaSEj8xXJm6AadFt26edQihPRdithY,25897
 tpu_inference/kernels/fused_moe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/kernels/fused_moe/v1/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/kernels/fused_moe/v1/kernel.py,sha256=xVXfclgbw_3U7c5W1azDFkFDK5FolBzDN9IL0rIzLQs,62813
+tpu_inference/kernels/fused_moe/v1/kernel.py,sha256=QHB0QEvC3x_6zhwz06JQpaOncQcNAhOSV92dD5tGVq8,40869
 tpu_inference/kernels/mla/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/kernels/mla/v1/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/kernels/mla/v1/kernel.py,sha256=dw1nhpL47uQxMFOIN2kENC6aITbalT81YZLAyr1usLU,51571
@@ -67,28 +66,27 @@ tpu_inference/kernels/ragged_paged_attention/v2/kernel.py,sha256=OiQGAHhyggbp1Pe
 tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py,sha256=vGp2ZWODTbjyG9z2z0Qf_BX-wYHd5bUybnc_DtOz0nI,10995
 tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py,sha256=mw80bXBGenroGdrITV0F_EaI2s-Z9KWwqU9WodvJg14,97919
 tpu_inference/kernels/ragged_paged_attention/v3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/kernels/ragged_paged_attention/v3/kernel.py,sha256=O179Fft5KpuN5LIFx3SghWXJJUqh3Og-xqfO4Z8QXYU,57032
-tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py,sha256=1ysmx7awuSZUnR7TcyUkARAvyMxNQS-9XRFMYnadZvk,61195
-tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py,sha256=k3LwduhZO85cJ-pSgnGN0c2Nn8eNeQq4eA94KUXJzMw,142198
+tpu_inference/kernels/ragged_paged_attention/v3/kernel.py,sha256=tlP6121yfXaukx_RQroHlHcZnbKPyyum0lAcvT0B_Pk,56132
+tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py,sha256=DFVdIIKmyufu_4b-3YhxI56jt0O1cJ3JsVl-2DDZHv4,55350
+tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py,sha256=leTS75aq99N1Zuv6wB5yLdkfYnEtrBDVI4z_jOKnjL0,142012
 tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes_hd64.py,sha256=P3_ivi8iUz5QMU_3pgpl4Bkbmn0q0NpDtVJX39haRQA,11208
 tpu_inference/kernels/ragged_paged_attention/v3/util.py,sha256=1N_ozjKboDYLteFJndWoLXNudj2z53rGXMkELa5Z9tY,1102
 tpu_inference/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/layers/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/layers/common/attention_interface.py,sha256=SQZ-1I32Jqg7GGI-z4BVibXbaitJHyTs26X3B5nBRVo,13369
 tpu_inference/layers/common/attention_metadata.py,sha256=St8ZatbY1D7xQACKJH459jMgp3oTP3AQ36mi9FZdrPU,850
-tpu_inference/layers/common/binary_search.py,sha256=ZQi-z1wG6WTcfVQXeTGOZokX4K1DSf9kCzqfrhEU8lk,12320
-tpu_inference/layers/common/quant_methods.py,sha256=mQSxZ44-QQtm22C_8ViejnP1cP2Dv6yc2YaP6oMKJeQ,185
-tpu_inference/layers/common/sharding.py,sha256=sjbwkDr2fP26Ob8f5cSDeDifr3eWFZMDHU4MKr7pIgQ,25217
 tpu_inference/layers/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/jax/attention_interface.py,sha256=1jlvSZWaP6DuPVtb1W_KPw4-Qi68BikOBNLLcpygupY,13221
 tpu_inference/layers/jax/base.py,sha256=Vhts6ZMwNCZ8LbnEXeB0rl3nHdS5hDJWX7HEa7Fl7yE,5775
+tpu_inference/layers/jax/binary_search.py,sha256=ZQi-z1wG6WTcfVQXeTGOZokX4K1DSf9kCzqfrhEU8lk,12320
 tpu_inference/layers/jax/constants.py,sha256=NcYg0zAf3ClfP7YMYdYu_F1GngOzZaIxIAHBZDunKw4,2755
 tpu_inference/layers/jax/layers.py,sha256=yv_lC2tbJuzVL-OaXYooX82Ys8hWZATeH9M78coJ3VI,10633
 tpu_inference/layers/jax/misc.py,sha256=znKv1Nuq_LgYpaIu0qlzUVDgQWnjjG7aqPJGM8kuwcw,566
 tpu_inference/layers/jax/rope.py,sha256=i2E7pRLWgOaFLbeo8_phZwKQWJW7ohAyl69E2V2Mc2U,11349
 tpu_inference/layers/jax/rope_interface.py,sha256=X0SruXizlCHGnssFujC1pL07UC4Vsp7-gdBy_Q7JZhI,8375
+tpu_inference/layers/jax/sharding.py,sha256=wBqdkXZSWfnnH8pkJtyW2DSqmAe_V4Vxi0iMPaXq0Z0,25185
 tpu_inference/layers/jax/transformer_block.py,sha256=ufv-yfVDmRP_Ynrx3UX9xj-x0PkNw_tQ-0N0eYf4i7M,3917
 tpu_inference/layers/jax/attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/layers/jax/attention/attention.py,sha256=DJFDkpQc9SDD156wVPFw3r2XaBgb44QNJ8OcdONaF5g,10085
+tpu_inference/layers/jax/attention/attention.py,sha256=bWXMtF8TToiRyZ3SvJGQnD2urZTbuX_omHVXpQdn0fs,10082
 tpu_inference/layers/jax/attention/deepseek_v3_attention.py,sha256=YlagoBMwINv2KRH1dr4oEcH_cQ9QMPB55nO2FQZsWs0,14010
 tpu_inference/layers/jax/attention/gpt_oss_attention.py,sha256=rkrEv4aNZxtAGcXd1HXHUxhNeDNAd9nWTEZOKWSI8cA,8725
 tpu_inference/layers/jax/attention/llama4_attention.py,sha256=VvUmfBxQEbHf3F2BrcYDUnq5abj7CSDYeRsNx_eVAh0,6162
@@ -97,78 +95,82 @@ tpu_inference/layers/jax/moe/deepseek_v3_moe.py,sha256=Q6CuwwiZtWYm6iUee1wJoDJrw
 tpu_inference/layers/jax/moe/gpt_oss_moe.py,sha256=Rx5b1jg2XMm7Xx9hrjgvyhscaJ_zGbVMHmeEiLh7kIQ,6196
 tpu_inference/layers/jax/moe/moe.py,sha256=cA8R1rjbBwNEoNlsPWjeIBB9nvaRDwlEdwQTVg6lTpY,8762
 tpu_inference/layers/jax/sample/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/layers/jax/sample/rejection_sampler.py,sha256=nI5s0E73xkqDIu2hTljIXt23B1Q-gRnC1myoQpGDJrQ,20426
-tpu_inference/layers/jax/sample/sampling.py,sha256=C30KgmdOVSaagvHhbfLgVJtVQmJo86CbHPa4h36Vn70,3314
+tpu_inference/layers/jax/sample/rejection_sampler.py,sha256=IRfVWjkbVXp9Sv1YrGMMh-LYx1AwbY-3FTXEO1-Ue9g,20423
+tpu_inference/layers/jax/sample/sampling.py,sha256=dVOcMdmPdAEsupPk96tCaZecIWUiDej0DiVnwaH9ckQ,3308
 tpu_inference/layers/jax/sample/sampling_metadata.py,sha256=Gd835LNWfGM0NRQBVBqEv0nPwt5q9F4AdFym0CUS1fw,2561
 tpu_inference/layers/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/layers/vllm/attention.py,sha256=wbJpcgqEAuIirv5PIULbiP-ggMKjmTanbB7Dg0BVYv4,7366
-tpu_inference/layers/vllm/fused_moe.py,sha256=qGbQoCq-sdcZj_Q0kP6RzQk7_YvcX7FopkpLcerjNFM,17819
+tpu_inference/layers/vllm/attention.py,sha256=JxEQ8ql_97zbQzukIbfUYq50-2k81VUG1Km_YV_RUtg,7363
+tpu_inference/layers/vllm/fused_moe.py,sha256=XZt2CPUz00qZzDcyfBFz6buhVzmGL1amHalHJALl9zw,18945
 tpu_inference/layers/vllm/linear_common.py,sha256=_YlJtbdaYcck_j-gFLos_k0ycktVWxT8Qo57tR2YqJ8,7749
-tpu_inference/layers/vllm/sharding.py,sha256=as7CF8UKTF3ToymwRY5Pi8uzwJk0P1sHPkWB5xEx3mA,9169
-tpu_inference/layers/vllm/quantization/__init__.py,sha256=SEppGayBzzQ5tsXLSy99aqilkAawQwYxnv2alCg6-ZU,1777
-tpu_inference/layers/vllm/quantization/awq.py,sha256=-8ZmjGvSKJB6_JuwSctNWt8xHWq4VSvK_AK9iahlgCo,8495
-tpu_inference/layers/vllm/quantization/common.py,sha256=8XD64pPa077c9HThFhLFVHlDL9YBafnYwp6rp6gR44E,4432
-tpu_inference/layers/vllm/quantization/mxfp4.py,sha256=UT6gpMrH27CusdGUMqEvQpJg1CPvsvnqAe0GKfZdV6o,13596
-tpu_inference/layers/vllm/quantization/unquantized.py,sha256=YaZdO_XjT06U1gtsUgNVSF1BrFqc4sCGO0dgtprUtwM,14395
+tpu_inference/layers/vllm/sharding.py,sha256=WTx1tF_7R99AdyE-lL7HQJ378hAafeI-JVRsugAvwn4,9177
+tpu_inference/layers/vllm/quantization/__init__.py,sha256=Tz44kUZTdNFu5Dmu48aQ-9f7ioWjbUWS0eVYURXZ17E,1535
+tpu_inference/layers/vllm/quantization/awq.py,sha256=ar8x1CPTPvfcf4wbuBC1XVh4pjtSUchoYWnbkZKH3CQ,8412
+tpu_inference/layers/vllm/quantization/common.py,sha256=wm3pge6XMTMsLK7_SSdgBP0PvQzz-1mrqN2I6xMqzrc,4218
+tpu_inference/layers/vllm/quantization/unquantized.py,sha256=id6d_IZIhDIvmaH3ANtmLiy4U_uY_AYAf4KTvfs3nmc,14900
 tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py,sha256=6idEyy3e849fZ1UeNvc9eSHYX7e6qvohrJa_d_D9MBk,5285
+tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py,sha256=uKaauZhaRDcMqd8_NyQoFs9BazMOFix3nIuutbLHHbU,5123
 tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors_moe.py,sha256=FM901QhyhJRC8CuMeICzCVVERvBHbhruRxYW0EQ570s,8820
 tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py,sha256=6sQvsxiWdi5Vte8V9vrQ2abaqGqWpq-mtzU7lGAo-ac,8759
 tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py,sha256=4y7lYgybpXszpCAtxGFhR8LDEbEoCCeo3DfUSOXxhaQ,5202
 tpu_inference/lora/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/lora/torch_lora_ops.py,sha256=pr3N7DVfkn3ANijUC6dBoiCtIJW4fdJpKdC3zWBUsxE,3121
-tpu_inference/lora/torch_punica_tpu.py,sha256=qTnXZGLoOgvukSxeunO_SfpPTlkq9GlMj9H7zVYg9LE,12680
+tpu_inference/lora/torch_punica_tpu.py,sha256=b27DpmIS_N5bhlIcryiENYNmPxp_cu40CGxjPW64d44,12706
+tpu_inference/mock/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/mock/vllm_config_utils.py,sha256=FlQshLjoHdgs3C66tYHYbKFUjbk9DhUwY-7HibZk0fI,878
+tpu_inference/mock/vllm_envs.py,sha256=cCubeOhH2WeYZQFJt6W0y_IiQo0fzIWR1LCCE8i6kI4,50990
+tpu_inference/mock/vllm_logger.py,sha256=vUGnN5nKT--ZvU15YCzODUM_FGiXKhcrrjDGjeN00RQ,7297
+tpu_inference/mock/vllm_logging_utils.py,sha256=TEUmKj3xHiLzHBnFqAujcxH0t2hBQ04sUaho2RyORnk,486
 tpu_inference/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/models/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/models/common/model_loader.py,sha256=b3aigca81gMVJt42oF2aoRohQHjBBe3oK3IPblZAaUM,19996
+tpu_inference/models/common/model_loader.py,sha256=AwukmGaUq2wv3OnFHUU-nwdAnKLG_eGw7PYY5CNrNNI,18225
 tpu_inference/models/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/models/jax/deepseek_v3.py,sha256=SKOHVEC-_2NLxBnzBzbu5tu0d6FTlAEiI1EefGaO2QE,40047
 tpu_inference/models/jax/gpt_oss.py,sha256=Vw4LRB5Kp6hbA2hjZGFS8kiEqOCjf881XH2JNtu2S1I,20924
 tpu_inference/models/jax/jax_intermediate_tensor.py,sha256=Pxu1PCV5LN5X58aYVkPiohcXZIeKVim2oqvrS_cVgw4,2604
-tpu_inference/models/jax/llama3.py,sha256=ZiFtrpAzXTT9vAPES9UeuJInCWGbvDWs7g0_JLdCCa4,13479
+tpu_inference/models/jax/llama3.py,sha256=YUG0S0Y6cy7PLcq0cpmDsGWbOZIhZzzyObRQdmUUxkg,13420
 tpu_inference/models/jax/llama4.py,sha256=wf2Sp2iYViaYD5rSfv3_ryO6gYuYM5XaOyvghaP4OCY,29631
-tpu_inference/models/jax/llama_eagle3.py,sha256=7-U99yvBkle-FSZ3NDDI-obWSQ2Fo2OTOi1H67H4jxY,12476
-tpu_inference/models/jax/llama_guard_4.py,sha256=LrnU2zBWM0s4q_5dwmR--OO0V7ttltsYhrHYlBgQVIw,15275
-tpu_inference/models/jax/qwen2.py,sha256=SuAp7tErk8OoIRko0Vt6QSOZP_9B9r5GTfqmVfImUIo,13410
-tpu_inference/models/jax/qwen2_5_vl.py,sha256=WUOmqNE6fHQ8PGU85Y8Bt6-CtCC1Uubbox_9FdpDMMo,49833
-tpu_inference/models/jax/qwen3.py,sha256=CIZQKjZDke_LPGsLNhRCJdDTzWueUneBPAQ1blS24IM,11050
+tpu_inference/models/jax/llama_eagle3.py,sha256=STUkAK6XEA7JM3i_Lx36-t5BhkAGeW_xYiq3zYhHP1A,12297
+tpu_inference/models/jax/phi3.py,sha256=Oz68PE2Z1t8wTed95_w0KMIXfnfV72ZwXugNOdWOV5w,13576
+tpu_inference/models/jax/qwen2.py,sha256=RYb0hMKzPnFOAyhqbztoNlSrFIlRa74fYqSNecA2VOY,13354
+tpu_inference/models/jax/qwen2_5_vl.py,sha256=J4-AjeS_igJdxYCjTwS0HShiEfwQUMwrHxjlWvMw0ok,43939
+tpu_inference/models/jax/qwen3.py,sha256=SOL-Pvp56IrMxqXpPf5EFacBI6AJNlqf4Zrr1pkabGw,10994
 tpu_inference/models/jax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/models/jax/utils/file_utils.py,sha256=NOuSC3YFnZpf3CZgYdghbbiNYJt42zgjlEYbOZIVct4,2840
 tpu_inference/models/jax/utils/multi_modal_utils.py,sha256=rrIrQWidkUnGilBHKNpdYh7_2BkvnAaqanXjC81GNcg,6156
-tpu_inference/models/jax/utils/weight_utils.py,sha256=qFU53jPHPvIcs_EOdIH80oNojpUp7GdSY2E6NZNsjvM,21376
+tpu_inference/models/jax/utils/weight_utils.py,sha256=65-H8BTbyilIBMBfvWjkkW3mf4soYASbhrJFqbFKzL4,20129
 tpu_inference/models/jax/utils/quantization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/models/jax/utils/quantization/mxfp4_utils.py,sha256=boGnqJCRIOf5nedAxQ8_IUTV6Rfll10DXnRC40BeeE8,3682
-tpu_inference/models/jax/utils/quantization/quantization_utils.py,sha256=rzAFU3OtQvg8w8ow0V15rMljAsa4SBrwOye6OI8Bty4,26530
+tpu_inference/models/jax/utils/quantization/quantization_utils.py,sha256=xgKoKB7AM3TYPxzVgEGLTK9ebQH2Kx8mNuO0heovkmk,26778
 tpu_inference/models/jax/utils/quantization/configs/fp8_all_modules_w_only.yaml,sha256=d_YHPtaRJ_7PBrPijSzJGnVeoJO62tKIGqrgFqpYT1k,137
 tpu_inference/models/jax/utils/quantization/configs/fp8_default.yaml,sha256=b7SyL75HuSTj3fN9_ZLCK_CDiccL5DGq_DddGmxj_qk,170
 tpu_inference/models/jax/utils/quantization/configs/int8_all_modules_w_only.yaml,sha256=0Qwij71zj9k6rmrUNd8Q5df9YYfkoJ1ZkgMAHxQy81k,128
 tpu_inference/models/jax/utils/quantization/configs/int8_default.yaml,sha256=lGec0UwwxmNPNgKPSsTsCMSXNJjhw507KMtM2NsSCMw,152
 tpu_inference/models/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/models/vllm/vllm_model_wrapper.py,sha256=3EcaD_1vZuyAZBfDtm5u_qfCahQU28qR4rAUraNAFqs,12305
+tpu_inference/models/vllm/vllm_model_wrapper.py,sha256=ERxj-cm-pmYpT9eiL-E3OxeaQDEDrH_Vs0iUS9nCU9s,11424
 tpu_inference/models/vllm/vllm_model_wrapper_context.py,sha256=yxlJHPmRQIAwlb1MmHK3xfXokgIkJ-evNU4PgyoJUdg,1187
 tpu_inference/platforms/__init__.py,sha256=lQCrKddS_GcGpCbeogvz9zOZD1mQw5bBsiw8On46qFQ,74
-tpu_inference/platforms/tpu_platform.py,sha256=q_eACjDkJkmnrUrKQzfK6hyqGEf2OjWn16-JHXwWquY,10723
+tpu_inference/platforms/tpu_platform.py,sha256=bdo_zlRqrhccpaz6zOdH18cU8kq6tGKgR1xJJehsVrc,10131
 tpu_inference/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/runner/block_table.py,sha256=K3Ic8EgPM08d_C5nEN60mxoRydlaQWySAemf_8Q_qVw,4175
-tpu_inference/runner/compilation_manager.py,sha256=dU0Yk8f0LtRTBe2q0iB3xcMSRco_WPsj2wS6zZJ8WhY,40375
+tpu_inference/runner/compilation_manager.py,sha256=pJFFLkFVmhXukBIxGRUo-hrOqx8jl8JIUuS36fZ2yvg,36177
 tpu_inference/runner/input_batch.py,sha256=bx221NX2IOWzrtopss-B-2ZKW4y-U6nQpG09PjpUziw,18273
-tpu_inference/runner/kv_cache.py,sha256=LKOZM5o8_62KDXhhYzQl2ibifgxN89ZxHvB1NT9u3MQ,4577
-tpu_inference/runner/kv_cache_manager.py,sha256=N0a896CE7Zrs_d4ZSSzRdqgjV1It57RBDSIpOzkRqro,22013
+tpu_inference/runner/kv_cache.py,sha256=i54EbGQB-9bbOgk6KibTpJpTE2pfFuTfis7J1P_UB0M,4574
+tpu_inference/runner/kv_cache_manager.py,sha256=CJxXtdWuewJqcTBMoR70_Uvwxjtc3cK2jxe1KpI9kQc,22152
 tpu_inference/runner/lora_utils.py,sha256=B4xMCgXGJ4VNdePvn89HH3tIZ-gYsQ7Vq_YCiYIATEY,3843
 tpu_inference/runner/multimodal_manager.py,sha256=azEPdHOwz8CN11MQmorGdtrCLbFaTCxdWyuEsZTzjYM,9778
-tpu_inference/runner/persistent_batch_manager.py,sha256=Otu67vOTf1_HKAMZgPDDHlRvvZ3YVJdz-QderH4qOII,13263
+tpu_inference/runner/persistent_batch_manager.py,sha256=KERSfKy6XjMejnbtPGI3hzoYAHJLeCxmpZVYPqBCago,11156
 tpu_inference/runner/speculative_decoding_manager.py,sha256=I3FDWKh2dn6nV8LgTGfCTwMKYnxQsTPpBIrmaJngXHs,10215
-tpu_inference/runner/structured_decoding_manager.py,sha256=gZQKQUFxh6xYYH9eGTdbguqk8hc2WwTrIdMMuCcbymE,3573
-tpu_inference/runner/tpu_runner.py,sha256=NBDKfSGShHmYpudrtGfo1hnVSQTcLpZV_nPiXEo7JPQ,79439
-tpu_inference/runner/utils.py,sha256=lKqL5nxGTk7ufzJRNdp4udn2bPu3jIX52W7akXgSrHc,17133
+tpu_inference/runner/structured_decoding_manager.py,sha256=Y0ERPhj4olFh6Y2TxP0R1_4UIJwy7nemYA-h63YIR2U,3622
+tpu_inference/runner/tpu_runner.py,sha256=5vPFey3KFnh5lczyj4cIT3mVhR8RuX8kbcuHVOg8DAg,72318
+tpu_inference/runner/utils.py,sha256=ZnWUoNo-7INeB0mdXti1jwUOdbmxyExznOs-crRTQLk,17126
 tpu_inference/spec_decode/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tpu_inference/spec_decode/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/spec_decode/jax/eagle3.py,sha256=FxP0uWeQlHlgCpt1nY3FUd4lKlegKJljHyc05jJucaQ,19104
+tpu_inference/spec_decode/jax/eagle3.py,sha256=A1dt-dmBttpy-5DGcL4noEDCB0OGP8Xo6MXqgJvWIo8,16593
 tpu_inference/worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tpu_inference/worker/tpu_worker.py,sha256=LnZcSNxdhh0NkoWXxS5bZ0bsTMduSANehy2wELAaVsY,20672
-tpu_inference-0.0.1rc1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-tpu_inference-0.0.1rc1.dist-info/METADATA,sha256=Ckyu7tcPAfxr698v8vDxUI70CyEVWLVDvUMFFcgqYYQ,5503
-tpu_inference-0.0.1rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tpu_inference-0.0.1rc1.dist-info/top_level.txt,sha256=gb1hRIQ3DOawUfVzvPL2E__2KPIl9I0vb5r0xcRBGYQ,20
-tpu_inference-0.0.1rc1.dist-info/RECORD,,
+tpu_inference/worker/tpu_worker.py,sha256=KY7fH--NP7jiTduP5m0gDnmB2LbhIel0Ts37XmjYpPM,14207
+tpu_inference-0.11.1.dev202511130813.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+tpu_inference-0.11.1.dev202511130813.dist-info/METADATA,sha256=LARdH4AAJfZrrU2Pj4EIN8Zl0QLjzEpzkRCqBbeUdT8,5465
+tpu_inference-0.11.1.dev202511130813.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tpu_inference-0.11.1.dev202511130813.dist-info/top_level.txt,sha256=gb1hRIQ3DOawUfVzvPL2E__2KPIl9I0vb5r0xcRBGYQ,20
+tpu_inference-0.11.1.dev202511130813.dist-info/RECORD,,

tests/test_envs.py

@@ -1,203 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the tpu-inference project
-
-import pytest
-
-import tpu_inference.envs as envs
-from tpu_inference.envs import enable_envs_cache, environment_variables
-
-
-def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
-    assert envs.JAX_PLATFORMS == ""
-    assert envs.PHASED_PROFILING_DIR == ""
-    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
-    monkeypatch.setenv("PHASED_PROFILING_DIR", "/tmp/profiling")
-    assert envs.JAX_PLATFORMS == "tpu"
-    assert envs.PHASED_PROFILING_DIR == "/tmp/profiling"
-
-    assert envs.TPU_NAME is None
-    assert envs.TPU_ACCELERATOR_TYPE is None
-    monkeypatch.setenv("TPU_NAME", "my-tpu")
-    monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v5litepod-16")
-    assert envs.TPU_NAME == "my-tpu"
-    assert envs.TPU_ACCELERATOR_TYPE == "v5litepod-16"
-
-    # __getattr__ is not decorated with functools.cache
-    assert not hasattr(envs.__getattr__, "cache_info")
-
-
-def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
-    monkeypatch.setenv("TPU_NAME", "my-tpu")
-
-    # __getattr__ is not decorated with functools.cache
-    assert not hasattr(envs.__getattr__, "cache_info")
-
-    enable_envs_cache()
-
-    # __getattr__ is decorated with functools.cache
-    assert hasattr(envs.__getattr__, "cache_info")
-    start_hits = envs.__getattr__.cache_info().hits
-
-    # 2 more hits due to JAX_PLATFORMS and TPU_NAME accesses
-    assert envs.JAX_PLATFORMS == "tpu"
-    assert envs.TPU_NAME == "my-tpu"
-    assert envs.__getattr__.cache_info().hits == start_hits + 2
-
-    # All environment variables are cached
-    for environment_variable in environment_variables:
-        envs.__getattr__(environment_variable)
-    assert envs.__getattr__.cache_info(
-    ).hits == start_hits + 2 + len(environment_variables)
-
-    # Reset envs.__getattr__ back to non-cached version to
-    # avoid affecting other tests
-    envs.__getattr__ = envs.__getattr__.__wrapped__
-
-
-def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
-    # Ensure clean environment for boolean vars by setting to default "0"
-    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
-    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "0")
-    monkeypatch.setenv("NEW_MODEL_DESIGN", "0")
-    monkeypatch.setenv("USE_MOE_EP_KERNEL", "0")
-
-    # Test SKIP_JAX_PRECOMPILE (default False)
-    assert envs.SKIP_JAX_PRECOMPILE is False
-    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "1")
-    assert envs.SKIP_JAX_PRECOMPILE is True
-    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
-    assert envs.SKIP_JAX_PRECOMPILE is False
-
-    # Test VLLM_XLA_CHECK_RECOMPILATION (default False)
-    assert envs.VLLM_XLA_CHECK_RECOMPILATION is False
-    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "1")
-    assert envs.VLLM_XLA_CHECK_RECOMPILATION is True
-    monkeypatch.setenv("VLLM_XLA_CHECK_RECOMPILATION", "0")
-    assert envs.VLLM_XLA_CHECK_RECOMPILATION is False
-
-    # Test NEW_MODEL_DESIGN (default False)
-    assert envs.NEW_MODEL_DESIGN is False
-    monkeypatch.setenv("NEW_MODEL_DESIGN", "1")
-    assert envs.NEW_MODEL_DESIGN is True
-
-    # Test USE_MOE_EP_KERNEL (default False)
-    assert envs.USE_MOE_EP_KERNEL is False
-    monkeypatch.setenv("USE_MOE_EP_KERNEL", "1")
-    assert envs.USE_MOE_EP_KERNEL is True
-
-
-def test_integer_env_vars(monkeypatch: pytest.MonkeyPatch):
-    # Ensure clean environment for integer vars by setting to defaults
-    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "1")
-    monkeypatch.setenv("NUM_SLICES", "1")
-
-    assert envs.PYTHON_TRACER_LEVEL == 1
-    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "3")
-    assert envs.PYTHON_TRACER_LEVEL == 3
-    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "0")
-    assert envs.PYTHON_TRACER_LEVEL == 0
-
-    # Test NUM_SLICES (default 1)
-    assert envs.NUM_SLICES == 1
-    monkeypatch.setenv("NUM_SLICES", "2")
-    assert envs.NUM_SLICES == 2
-    monkeypatch.setenv("NUM_SLICES", "4")
-    assert envs.NUM_SLICES == 4
-
-
-def test_model_impl_type_choices(monkeypatch: pytest.MonkeyPatch):
-    # Test case sensitive choices
-    monkeypatch.setenv("MODEL_IMPL_TYPE", "flax_nnx")
-    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
-
-    monkeypatch.setenv("MODEL_IMPL_TYPE", "vllm")
-    assert envs.MODEL_IMPL_TYPE == "vllm"
-
-
-def test_string_env_vars_defaults(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.delenv("JAX_PLATFORMS", raising=False)
-    monkeypatch.delenv("PREFILL_SLICES", raising=False)
-    monkeypatch.delenv("DECODE_SLICES", raising=False)
-
-    assert envs.JAX_PLATFORMS == ""
-    assert envs.PREFILL_SLICES == ""
-    assert envs.DECODE_SLICES == ""
-    assert envs.PHASED_PROFILING_DIR == ""
-
-
-def test_none_default_env_vars(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.delenv("TPU_ACCELERATOR_TYPE", raising=False)
-    monkeypatch.delenv("TPU_NAME", raising=False)
-    monkeypatch.delenv("TPU_WORKER_ID", raising=False)
-
-    assert envs.TPU_ACCELERATOR_TYPE is None
-    assert envs.TPU_NAME is None
-    assert envs.TPU_WORKER_ID is None
-
-
-def test_ray_env_vars(monkeypatch: pytest.MonkeyPatch):
-    assert envs.RAY_USAGE_STATS_ENABLED == "0"
-    monkeypatch.setenv("RAY_USAGE_STATS_ENABLED", "1")
-    assert envs.RAY_USAGE_STATS_ENABLED == "1"
-
-    assert envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "shm"
-
-
-def test_invalid_attribute_raises_error():
-    with pytest.raises(AttributeError,
-                       match="has no attribute 'NONEXISTENT_VAR'"):
-        _ = envs.NONEXISTENT_VAR
-
-
-def test_dir_returns_all_env_vars():
-    env_vars = envs.__dir__()
-    assert isinstance(env_vars, list)
-    assert len(env_vars) == len(environment_variables)
-    assert "JAX_PLATFORMS" in env_vars
-    assert "TPU_NAME" in env_vars
-    assert "SKIP_JAX_PRECOMPILE" in env_vars
-    assert "VLLM_XLA_CHECK_RECOMPILATION" in env_vars
-    assert "MODEL_IMPL_TYPE" in env_vars
-
-
-def test_tpu_multihost_env_vars(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("TPU_WORKER_ID", "0")
-    assert envs.TPU_WORKER_ID == "0"
-
-    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "ray")
-    assert envs.TPU_MULTIHOST_BACKEND == "ray"
-
-
-def test_disaggregated_serving_env_vars(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("PREFILL_SLICES", "0,1,2,3")
-    assert envs.PREFILL_SLICES == "0,1,2,3"
-
-    monkeypatch.setenv("DECODE_SLICES", "4,5,6,7")
-    assert envs.DECODE_SLICES == "4,5,6,7"
-
-
-def test_model_impl_type_default(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.delenv("MODEL_IMPL_TYPE", raising=False)
-    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
-
-
-def test_cache_preserves_values_across_env_changes(
-        monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
-
-    enable_envs_cache()
-
-    assert envs.JAX_PLATFORMS == "tpu"
-
-    # Change environment variable
-    monkeypatch.setenv("JAX_PLATFORMS", "cpu")
-
-    # Cached value should still be "tpu"
-    assert envs.JAX_PLATFORMS == "tpu"
-
-    # Reset envs.__getattr__ back to non-cached version
-    envs.__getattr__ = envs.__getattr__.__wrapped__
-
-    # Now it should reflect the new value
-    assert envs.JAX_PLATFORMS == "cpu"

tpu_inference/layers/common/quant_methods.py

@@ -1,8 +0,0 @@
-UNQUANTIZED = "unquantized"
-MXFP4 = "mxfp4"
-AWQ = "awq"
-COMPRESSED_TENSORS = "compressed-tensors"
-
-
-def get_tpu_quant_method(quant_method: str) -> str:
-    return "tpu-" + quant_method