tpu-inference 0.11.1.dev202511150811__py3-none-any.whl → 0.11.1.dev202511270815__py3-none-any.whl

tpu_inference/runner/compilation_manager.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -135,12 +135,6 @@ class CompilationManager:
                 ShardingAxisName.ATTN_DATA, )) if dp_size > 1 else None
 
         # Keep existing pattern for complex array operations
-        block_tables = self.runner.block_table_cpu[:self.runner.max_num_reqs]
-        block_tables = block_tables.reshape(-1)
-        block_tables = device_array(self.runner.mesh,
-                                    block_tables,
-                                    sharding=dp_sharding)
-
         seq_lens = self._create_dummy_tensor((self.runner.max_num_reqs, ),
                                              jnp.int32, dp_sharding)
         query_start_loc = self._create_dummy_tensor(
@@ -152,26 +146,45 @@ class CompilationManager:
                                             request_distribution,
                                             sharding=dp_sharding)
 
-        attention_metadata = AttentionMetadata(
-            input_positions=positions,
-            block_tables=block_tables,
-            seq_lens=seq_lens,
-            query_start_loc=query_start_loc,
-            request_distribution=request_distribution,
-        )
+        attention_metadata_per_layer: Dict[str, AttentionMetadata] = {}
+        uniform_attention_metadata: AttentionMetadata = None
+        for kv_cache_gid, kv_cache_group in enumerate(
+                self.runner.kv_cache_config.kv_cache_groups):
+            block_tables = self.runner.block_tables_cpu[
+                kv_cache_gid][:self.runner.max_num_reqs]
+            block_tables = block_tables.reshape(-1)
+            block_tables = device_array(self.runner.mesh,
+                                        block_tables,
+                                        sharding=dp_sharding)
+
+            attention_metadata_gid = AttentionMetadata(
+                input_positions=positions,
+                block_tables=block_tables,
+                seq_lens=seq_lens,
+                query_start_loc=query_start_loc,
+                request_distribution=request_distribution,
+            )
+            if not self.runner.use_hybrid_kvcache:
+                # all layers share the same attention metadata
+                uniform_attention_metadata = attention_metadata_gid
+            else:
+                for layer_name in kv_cache_group.layer_names:
+                    attention_metadata_per_layer[
+                        layer_name] = attention_metadata_gid
 
         def model_fn_wrapper(
             state,
             kv_caches,
             input_ids,
             attention_metadata,
+            positions,
             inputs_embeds,
             layer_name_to_kvcache_index,
             lora_metadata,
         ):
             kv_caches, hidden_states, _ = self.runner.model_fn(
                 state, kv_caches, input_ids, attention_metadata, inputs_embeds,
-                layer_name_to_kvcache_index, lora_metadata)
+                positions, layer_name_to_kvcache_index, lora_metadata)
             self.runner.kv_caches = kv_caches
             return hidden_states
 
@@ -179,6 +192,10 @@ class CompilationManager:
                 self.runner.lora_config, np.array([num_tokens],
                                                   dtype=np.int32)):
             lora_metadata = self.runner.lora_utils.extract_lora_metadata()
+            if self.runner.use_hybrid_kvcache:
+                attention_metadata = attention_metadata_per_layer
+            else:
+                attention_metadata = uniform_attention_metadata
             self._run_compilation(
                 name,
                 model_fn_wrapper,
@@ -186,6 +203,7 @@ class CompilationManager:
                 self.runner.kv_caches,
                 input_ids,
                 attention_metadata,
+                positions,
                 inputs_embeds,
                 tuple(self.runner.layer_name_to_kvcache_index.items()),
                 lora_metadata,
@@ -332,13 +350,15 @@ class CompilationManager:
             index_paddings = self.runner.num_reqs_paddings
         dp_sharding = NamedSharding(self.runner.mesh,
                                     PartitionSpec(ShardingAxisName.ATTN_DATA))
+        hidden_states_sharding = NamedSharding(
+            self.runner.mesh, PartitionSpec(ShardingAxisName.ATTN_DATA, None))
         dp_size = self.runner.vllm_config.sharding_config.total_dp_size
         self._precompile_select_from_array_helper(
             name="select all logits",
             source_paddings=self.runner.num_tokens_paddings,
             indices_paddings=index_paddings,
             hidden_dim=hsize,
-            input_sharding=dp_sharding,
+            input_sharding=hidden_states_sharding,
             indices_sharding=dp_sharding if dp_size > 1 else None,
         )
 
@@ -528,7 +548,9 @@ class CompilationManager:
     def _precompile_eagle3_helpers(self) -> None:
         logger.info(
             "Compiling eagle3 jitted helpers with different input shapes.")
-        hidden_size = self.runner.model_config.get_hidden_size()
+        target_hidden_size = self.runner.model_config.get_hidden_size()
+        draft_hidden_size = self.runner.speculative_config.draft_model_config.get_hidden_size(
+        )
         dtype = self.runner.model_config.dtype
 
         num_kv_cache_groups = len(self.runner.kv_cache_config.kv_cache_groups)
@@ -575,7 +597,7 @@ class CompilationManager:
 
         for num_logits in self.runner.num_logits_paddings:
             hidden_states = self._create_dummy_tensor(
-                (num_logits, hidden_size), jnp.bfloat16)
+                (num_logits, draft_hidden_size), jnp.bfloat16)
             self._run_compilation(
                 "eagle3_get_draft_token_ids",
                 self.runner.drafter._get_draft_token_ids,
@@ -586,8 +608,8 @@ class CompilationManager:
         input_ids_loop = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32,
             NamedSharding(self.runner.mesh, PartitionSpec()))
-        target_hidden_state_loop = self._create_dummy_tensor(
-            (self.runner.max_num_reqs, hidden_size), dtype,
+        draft_hidden_state_loop = self._create_dummy_tensor(
+            (self.runner.max_num_reqs, draft_hidden_size), dtype,
             NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
         next_token_ids = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
@@ -595,9 +617,12 @@ class CompilationManager:
             (self.runner.max_num_reqs, ), jnp.int32)
         for num_tokens in self.runner.num_tokens_paddings:
             aux_hidden_states = [
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
-                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
+                self._create_dummy_tensor((num_tokens, target_hidden_size),
+                                          dtype),
             ]
 
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32)
@@ -628,15 +653,15 @@ class CompilationManager:
             input_ids = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             aux_hidden_states = [
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16,
+                    (num_tokens, target_hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
             ]
@@ -668,17 +693,17 @@ class CompilationManager:
                 state,
                 kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
             ):
                 kv_caches, hidden_states, _ = self.runner.drafter.model_fn(
-                    state, kv_caches, input_ids, target_hidden_states,
+                    state, kv_caches, input_ids, draft_hidden_states,
                     attention_metadata)
                 self.runner.kv_caches = kv_caches
                 return hidden_states
 
-            target_hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), dtype,
+            draft_hidden_states = self._create_dummy_tensor(
+                (num_tokens, draft_hidden_size), dtype,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, "model")))
             input_ids = self._create_dummy_tensor(
                 (num_tokens, ), jnp.int32,
@@ -689,7 +714,7 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids,
-                target_hidden_states,
+                draft_hidden_states,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
@@ -721,13 +746,13 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids_loop,
-                target_hidden_state_loop,
+                draft_hidden_state_loop,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
 
             hidden_states = self._create_dummy_tensor(
-                (num_tokens, hidden_size), jnp.bfloat16,
+                (num_tokens, draft_hidden_size), jnp.bfloat16,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
 
             self._run_compilation(

tpu_inference/runner/kv_cache_manager.py

@@ -1,15 +1,16 @@
 import functools
-import math
 from typing import TYPE_CHECKING, Dict, List
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 import vllm.envs as envs
 from jax.sharding import NamedSharding, PartitionSpec
 from torchax.ops.mappings import t2j_dtype
-from vllm.attention import Attention
 from vllm.attention.backends.abstract import AttentionType
+from vllm.attention.layer import Attention
 from vllm.config import get_layers_from_vllm_config
+from vllm.utils.math_utils import cdiv
 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
                                         KVCacheSpec, MLAAttentionSpec,
                                         SlidingWindowSpec)
@@ -175,6 +176,11 @@ class KVCacheManager:
             )
             self.runner.input_batch = new_input_batch
             self.runner.persistent_batch_manager.input_batch = new_input_batch
+            self.runner.block_tables_cpu = [
+                np.zeros((self.runner.max_num_reqs,
+                          cdiv(self.runner.max_model_len, block_size)),
+                         dtype=np.int32) for block_size in block_sizes
+            ]
 
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         self.maybe_reinitialize_input_batch(kv_cache_config)
@@ -190,7 +196,7 @@ class KVCacheManager:
             num_blocks = kv_cache_tensor.size // page_size_bytes
             dp_size = self.runner.vllm_config.sharding_config.total_dp_size
             # num_blocks must be a multiple of dp_size
-            num_blocks = math.ceil(num_blocks / dp_size) * dp_size
+            num_blocks = (num_blocks // dp_size) * dp_size
             # NOTE: we'll multiply the num_kv_heads by 2 in the function
             kv_cache = create_kv_caches(
                 num_blocks=num_blocks,

tpu_inference/runner/structured_decoding_manager.py

@@ -61,11 +61,10 @@ class StructuredDecodingManager:
         self.runner.require_structured_out_cpu.fill(0)
 
         sorted_struct_requests = sorted(
-            grammar_output.structured_output_request_ids.items(),
-            key=lambda item: item[1])
+            grammar_output.structured_output_request_ids)
 
         cumulative_mask_idx = 0
-        for req_id, _ in sorted_struct_requests:
+        for req_id in sorted_struct_requests:
             if req_id not in self.runner.input_batch.req_id_to_index:
                 continue
             batch_index = self.runner.input_batch.req_id_to_index[req_id]