[perf] xllm_fused_qknorm_rope.so compiled+wired, xllm_cache

vllm/model_executor/models/mamba_cache.py

@@ -10,11 +10,11 @@ class MambaCacheManager:
     def __init__(self, dtype, num_mamba_layers, max_batch_size,
                  conv_state_shape, temporal_state_shape):
 
-        conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
+        conv_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
                                  conv_state_shape,
                                  dtype=dtype,
                                  device="cuda")
-        temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
+        temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
                                      temporal_state_shape,
                                      dtype=dtype,
                                      device="cuda")
@@ -101,6 +101,8 @@ class MambaCacheManager:
             self._move_out_if_already_occupied(
                 index=destination_index,
                 all_occupied_indices=all_occupied_indices)
+            for cache_t in self.mamba_cache:
+                cache_t[:, destination_index].zero_()
             self.mamba_cache_indices_mapping[cur_rid] = {
                 seq_id: destination_index
             }
@@ -206,7 +208,10 @@ class MambaCacheManager:
                                    finished_seq_groups_req_ids: List[str]):
         for req_id in finished_seq_groups_req_ids:
             if req_id in self.mamba_cache_indices_mapping:
-                self.mamba_cache_indices_mapping.pop(req_id)
+                seq_mapping = self.mamba_cache_indices_mapping.pop(req_id)
+                for cache_idx in seq_mapping.values():
+                    for cache_t in self.mamba_cache:
+                        cache_t[:, cache_idx].zero_()
 
     def _first_free_index_in_mamba_cache(
             self, indices_range: Optional[List[int]] = None) -> int:
@@ -219,4 +224,4 @@ class MambaCacheManager:
             if i not in all_occupied_indices:
                 return i
         raise Exception("Couldn't find a free spot in the mamba cache! This"
-                        "should never happen")
+                        "should never happen")

vllm/sequence.py

@@ -119,6 +119,7 @@ class RequestMetrics:
     scheduler_time: Optional[float] = None
     model_forward_time: Optional[float] = None
     model_execute_time: Optional[float] = None
+    num_cached_tokens: Optional[int] = None
 
 
 class SequenceDataDelta(
@@ -527,6 +528,11 @@ class Sequence:
         self._last_output_token_ids_offset = output_len
 
         # Return new tokens
+        if num_new_tokens == 0:
+            # During chunked prefill steps with no output yet, num_new_tokens=0.
+            # Python's [-0:] == [0:] returns the ENTIRE list — guard against this.
+            return []
+
         if num_new_tokens == 1:
             # Optimization for single decode token case
             # (which is what we have most of the time)
@@ -935,6 +941,12 @@ class SequenceGroupMetadataDelta(
     computed_block_nums: Optional[List[int]] = None
     state: Optional[SequenceGroupState] = msgspec.field(
         default_factory=lambda: SequenceGroupState())
+    # BI100 hybrid prefix-cache actions. Fields are appended for msgspec wire
+    # compatibility with the pre-existing array-like structure.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
 
 class SequenceGroupMetadata(
@@ -1000,6 +1012,12 @@ class SequenceGroupMetadata(
     # Zero means speculative decoding is disabled for some reasons.
     # TODO: We should maintain this states out of the sequence group.
     num_speculative_tokens: Optional[int] = None
+    # BI100 hybrid prefix-cache actions. These are internal scheduler-to-worker
+    # metadata and never surface through the OpenAI API.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
     def __post_init__(self):
         if self.seq_data is not None and self.token_chunk_size is None:
@@ -1046,6 +1064,14 @@ class SequenceGroupMetadata(
         self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size
         self.do_sample = sequence_group_metadata_delta.do_sample
         self.is_prompt = sequence_group_metadata_delta.is_prompt
+        self.computed_block_nums = (
+            sequence_group_metadata_delta.computed_block_nums)
+        self.gdn_restore_key = sequence_group_metadata_delta.gdn_restore_key
+        self.gdn_capture_points = (
+            sequence_group_metadata_delta.gdn_capture_points)
+        self.gdn_evict_keys = sequence_group_metadata_delta.gdn_evict_keys
+        self.gdn_segment_offsets = (
+            sequence_group_metadata_delta.gdn_segment_offsets)
 
     def finish_step(self) -> None:
         assert self.state is not None

vllm/worker/worker.py

@@ -216,6 +216,13 @@ class Worker(LocalOrDistributedWorkerBase):
         """
         # Profile the memory usage of the model and get the maximum number of
         # cache blocks that can be allocated with the remaining free memory.
+        # PRD: skip profile_run when num_gpu_blocks_override is set
+        _ovr = getattr(self.cache_config, 'num_gpu_blocks_override', None)
+        if _ovr is not None and _ovr > 0:
+            logger.info("Skipping profile_run -- num_gpu_blocks_override=%d", _ovr)
+            _cbs = self.get_cache_block_size_bytes()
+            _cpu = self.cache_config.swap_space_bytes // _cbs if _cbs > 0 else 256
+            return int(_ovr), int(_cpu)
         torch.cuda.empty_cache()
 
         # Execute a forward pass with dummy inputs to profile the memory usage