[BUGFIX][Mamba][Qwen3.5] Zero freed SSM cache blocks on GPU

vllm/v1/core/block_pool.py

@@ -180,6 +180,10 @@ def __init__(
 
         self.metrics_collector = metrics_collector
 
+        # Block IDs allocated since the last drain. The worker uses these
+        # to zero the corresponding GPU memory before the blocks are used.
+        self.new_block_ids: list[int] = []
+
     def get_cached_block(
         self, block_hash: BlockHash, kv_cache_group_ids: list[int]
     ) -> list[KVCacheBlock] | None:
@@ -347,8 +351,15 @@ def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
                 block.ref_cnt += 1
                 if self.metrics_collector:
                     self.metrics_collector.on_block_allocated(block)
+        self.new_block_ids.extend(block.block_id for block in ret)
         return ret
 
+    def take_new_block_ids(self) -> list[int]:
+        """Drain and return block IDs allocated since the last call."""
+        ids = self.new_block_ids
+        self.new_block_ids = []
+        return ids
+
     def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
         """
         If a block is cached in `cached_block_hash_to_block`, we reset its hash

vllm/v1/core/kv_cache_manager.py

@@ -489,6 +489,10 @@ def create_kv_cache_blocks(
         # Only create new KVCacheBlocks for non-empty blocks
         return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
 
+    def take_new_block_ids(self) -> list[int]:
+        """Drain and return block IDs allocated since the last call."""
+        return self.block_pool.take_new_block_ids()
+
     def new_step_starts(self) -> None:
         """Called when a new step is started."""
         self.coordinator.new_step_starts()

vllm/v1/core/sched/output.py

@@ -238,6 +238,11 @@ class SchedulerOutput:
     # EC Cache Connector metadata
     ec_connector_metadata: ECConnectorMetadata | None = None
 
+    # Block IDs freshly allocated from the pool during this scheduling step.
+    # The worker zeros the corresponding GPU memory before the blocks are used,
+    # preventing stale NaN/data from corrupting attention or SSM computation.
+    new_block_ids_to_zero: list[int] | None = None
+
     @classmethod
     def make_empty(cls) -> "SchedulerOutput":
         return cls(

vllm/v1/core/sched/scheduler.py

@@ -871,6 +871,8 @@ def schedule(self) -> SchedulerOutput:
         self.prev_step_scheduled_req_ids.clear()
         self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
 
+        new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() or None
+
         scheduler_output = SchedulerOutput(
             scheduled_new_reqs=new_reqs_data,
             scheduled_cached_reqs=cached_reqs_data,
@@ -886,6 +888,7 @@ def schedule(self) -> SchedulerOutput:
             # the previous and the current steps.
             finished_req_ids=self.finished_req_ids,
             free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+            new_block_ids_to_zero=new_block_ids_to_zero,
         )
 
         # NOTE(Kuntai): this function is designed for multiple purposes:

vllm/v1/worker/gpu_model_runner.py

@@ -473,6 +473,9 @@ def __init__(
         # self.model: nn.Module  # Set after load_model
         # Initialize in initialize_kv_cache
         self.kv_caches: list[torch.Tensor] = []
+        # (raw_tensor, page_size_bytes) pairs for zeroing newly allocated blocks.
+        # Populated by _allocate_kv_cache_tensors.
+        self.kv_cache_raw_buffers: list[tuple[torch.Tensor, int]] = []
         # Initialize in initialize_kv_cache_tensors
         self.cross_layers_kv_cache: torch.Tensor | None = None
         self.cross_layers_attn_backend: type[AttentionBackend] | None = None
@@ -906,6 +909,13 @@ def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
                 decode_threshold=self.reorder_batch_threshold,
             )
 
+    def _zero_block_ids(self, block_ids: list[int]) -> None:
+        """Zero the raw KV cache memory for the given block IDs."""
+        for raw_tensor, page_size in self.kv_cache_raw_buffers:
+            for bid in block_ids:
+                start = bid * page_size
+                raw_tensor[start : start + page_size] = 0
+
     # Note: used for model runner override.
     def _init_device_properties(self) -> None:
         """Initialize attributes from torch.cuda.get_device_properties"""
@@ -939,6 +949,11 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         for req_id in scheduler_output.finished_req_ids:
             self.input_batch.remove_request(req_id)
 
+        # Zero GPU memory for freshly allocated cache blocks to prevent
+        # stale NaN/data from corrupting attention or SSM computation.
+        if scheduler_output.new_block_ids_to_zero:
+            self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
+
         # Free the cached encoder outputs.
         for mm_hash in scheduler_output.free_encoder_mm_hashes:
             self.encoder_cache.pop(mm_hash, None)
@@ -5815,12 +5830,15 @@ def _allocate_kv_cache_tensors(
             corresponding memory buffer for KV cache.
         """
         kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
+        num_blocks = kv_cache_config.num_blocks
         for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             tensor = torch.zeros(
                 kv_cache_tensor.size, dtype=torch.int8, device=self.device
             )
             for layer_name in kv_cache_tensor.shared_by:
                 kv_cache_raw_tensors[layer_name] = tensor
+            page_size = kv_cache_tensor.size // num_blocks
+            self.kv_cache_raw_buffers.append((tensor, page_size))
 
         layer_names = set()
         for group in kv_cache_config.kv_cache_groups: