tpu-inference 0.13.2rc1__tar.gz → 0.13.2rc3__tar.gz

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/MANIFEST.in

@@ -1,4 +1,3 @@
 include requirements.txt
-include requirements_v7x.txt
 include README.md
 include tpu_inference/models/jax/utils/quantization/configs/*.yaml

{tpu_inference-0.13.2rc1/tpu_inference.egg-info → tpu_inference-0.13.2rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tpu_inference
-Version: 0.13.2rc1
+Version: 0.13.2rc3
 Author: tpu_inference Contributors
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/setup.py

@@ -20,40 +20,26 @@ def get_requirements() -> List[str]:
             requirements = f.read().strip().split("\n")
         resolved_requirements = []
         for line in requirements:
-            if not line or line.startswith("#"):
-                continue
             if line.startswith("-r "):
                 resolved_requirements += _read_requirements(line.split()[1])
-            elif line.startswith(("-", "--")):
+            elif line.startswith("--"):
                 continue
             else:
                 resolved_requirements.append(line)
         return resolved_requirements
 
     try:
-        #requirements = _read_requirements("requirements_v7x.txt")
-
-        # For TPU v7x build
-        if os.getenv("IS_FOR_V7X", "false").lower() == "true":
-            print("Using requirements_v7x.txt")
-            requirements = _read_requirements("requirements_v7x.txt")
-            #requirements.extend(v7x_requirements)
-        else:
-            #For TPU v6e build
-            print("Using requirements.txt")
-            requirements = _read_requirements("requirements.txt")
-
+        requirements = _read_requirements("requirements.txt")
     except ValueError:
         print("Failed to read requirements.txt in vllm_tpu.")
     return requirements
 
 
 def get_version():
-    version = os.getenv("VLLM_VERSION_OVERRIDE", "0.0.0").strip()
-    if os.getenv("IS_FOR_V7X", "false").lower() == "true":
-        version = f"{version}.post7"
+    if env_version := os.getenv("VLLM_VERSION_OVERRIDE"):
+        return env_version
+    return "0.0.0"
 
-    return version
 
 setup(
     name="tpu_inference",

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/tests/layers/jax/test_qwix.py

@@ -832,7 +832,7 @@ class TestGetDefaultQwixQuantizationConfig(unittest.TestCase):
         # Patch the constants in the module where the function resides
         self.patchers = [
             patch(
-                "tpu_inference.models.jax.utils.qwix.qwix_utils.DEFAULT_DEEPSEEK_FP8_CONFIG",
+                "tpu_inference.models.jax.utils.qwix.qwix_utils.DEFAULT_DEEPSEEK_FP4_MLP_MOE_FP8_ATTN_CONFIG",
                 self.mock_deepseek_config),
             patch(
                 "tpu_inference.models.jax.utils.qwix.qwix_utils.DEFAULT_LLAMA4_FP8_CONFIG",

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/tpu_inference/models/jax/utils/qwix/qwix_utils.py

@@ -35,7 +35,7 @@ DEFAULT_NUM_TOKENS_FOR_MODEL_INPUTS = 512
 DEFAULT_MAX_NUM_SEQS_FOR_MODEL_INPUTS = 256
 DEFAULT_MAX_NUM_BLOCKS_PER_REQ = 16
 
-DEFAULT_DEEPSEEK_FP8_CONFIG = {
+DEFAULT_DEEPSEEK_FP4_MLP_MOE_FP8_ATTN_CONFIG = {
     "qwix": {
         "use_abstract_model":
         True,
@@ -452,7 +452,7 @@ def get_default_qwix_quantization_config(
     # NOTE (jacobplatin): we'll default to mixed FP8 (attention) + FP4 (MoE experts)
     # for DeepSeek
     if model_type == "deepseek_v3" and quant_method == "fp8":
-        config = copy.deepcopy(DEFAULT_DEEPSEEK_FP8_CONFIG)
+        config = copy.deepcopy(DEFAULT_DEEPSEEK_FP4_MLP_MOE_FP8_ATTN_CONFIG)
 
         # Dynamically fetch block size from HF config if available
         # Config fmt: 'weight_block_size': [1, 512] -> we want the 2nd dim for tile_size
@@ -462,7 +462,7 @@ def get_default_qwix_quantization_config(
         block_size = hf_quant_config["weight_block_size"]
         if isinstance(block_size, (list, tuple)) and len(block_size) == 2:
             assert block_size[
-                0] == 1, f"Expected first dimension to be 1 (unchanneled), but got {block_size[0]}!"
+                0] == 1, f"Expected first dimension to be 1 (unchanneled), but got {block_size[0]}! If you are trying to run quantized DeepSeek, we currently only support 1D-subchannel quantization and those models can be found here: https://huggingface.co/collections/jrplatin/deepseek-r1-1d-subchannel"
             tile_size = block_size[1]
             assert tile_size > 1, f"Expected tile_size > 1 for DeepSeek, but got {tile_size}"
             logger.info(

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/tpu_inference/runner/kv_cache_manager.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, List
 
 import jax
 import jax.numpy as jnp
@@ -212,7 +212,6 @@ class KVCacheManager:
         # uniform page size.
         representative_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
         page_size_bytes = representative_spec.page_size_bytes
-        self.runner.layer_name_to_kvcache_index: Dict[str, int] = {}
         kv_caches = self.runner.kv_caches
         num_blocks_list = []
         for i, kv_cache_tensor in enumerate(kv_cache_config.kv_cache_tensors):

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/tpu_inference/runner/tpu_runner.py

@@ -282,6 +282,9 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         self._substitute_placeholder_token_fn = _substitute_placeholder_token
         self.execute_model_state: ExecuteModelState | None = None
 
+        self.kv_caches: list[jax.Array] = []
+        self.layer_name_to_kvcache_index: dict[str, int] = {}
+
     def _init_random(self):
         if self.model_config.seed is None:
             self.model_config.seed = 0
@@ -545,7 +548,6 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         self.topology_order_id = topology_order_id
         self.kv_cache_config = kv_cache_config
         self.use_hybrid_kvcache = len(kv_cache_config.kv_cache_groups) > 1
-        self.kv_caches = []
         self.kv_cache_manager.initialize_kv_cache(kv_cache_config)
         if has_kv_transfer_group():
             get_kv_transfer_group().register_runner(self)

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3/tpu_inference.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tpu_inference
-Version: 0.13.2rc1
+Version: 0.13.2rc3
 Author: tpu_inference Contributors
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers

{tpu_inference-0.13.2rc1 → tpu_inference-0.13.2rc3}/tpu_inference.egg-info/SOURCES.txt

@@ -3,7 +3,6 @@ MANIFEST.in
 README.md
 pyproject.toml
 requirements.txt
-requirements_v7x.txt
 setup.py
 tests/__init__.py
 tests/test_base.py

tpu_inference-0.13.2rc1/requirements_v7x.txt

@@ -1,25 +0,0 @@
-# This file contains additional dependencies needed for TPU v7x support.
-# It is expected to be used in conjunction with the main requirements.txt file.
---pre
--i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
--f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-jax==0.8.1
-jaxlib==0.8.1
-jaxtyping==0.3.2
-libtpu==0.0.31
-
-tpu-info==0.7.1
-yapf==0.43.0
-pytest
-pytest-mock
-absl-py
-numpy
-google-cloud-storage
-flax==0.11.1
-torchax==0.0.10
-qwix==0.1.1
-torchvision==0.24.0
-pathwaysutils
-parameterized
-numba==0.62.1
-runai-model-streamer[s3,gcs]==0.15.0