[v1][kvconnector] DecodeBenchConnector: fill list/tuple (Mamba/KDA) KV caches (#45080)

Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py

@@ -387,40 +387,28 @@ class DecodeBenchConnectorWorker:
 
             kv_cache = self.kv_caches[layer_name]
 
-            # Convert block_ids to tensor on device
-            block_ids_tensor = torch.tensor(
-                block_ids, dtype=torch.long, device=kv_cache.device
-            )
-
-            # Filter invalid block IDs
-            valid_mask = block_ids_tensor < kv_cache.shape[0]
-            valid_block_ids = block_ids_tensor[valid_mask]
-
-            if len(valid_block_ids) == 0:
-                continue
-
-            # Create fill values - either constant or random
-            block_shape = kv_cache.shape[1:]
-            if self.fill_std > 0:
-                # Random normal sampling
-                fill_values = torch.normal(
-                    mean=self.fill_mean,
-                    std=self.fill_std,
-                    size=(len(valid_block_ids),) + block_shape,
-                    dtype=kv_cache.dtype,
-                    device=kv_cache.device,
-                )
+            # Attention layers store KV as a single block-indexed tensor whose
+            # first dim is num_blocks; fill the requested block rows. Hybrid /
+            # linear-attention layers (e.g. Mamba, Kimi Delta Attention) store
+            # their state as a list/tuple of tensors that are NOT block-indexed
+            # — each tensor is a single state buffer with no num_blocks
+            # dimension — so fill each tensor in its entirety with the same
+            # dummy values.
+            if isinstance(kv_cache, torch.Tensor):
+                self._fill_block_tensor(kv_cache, block_ids)
+            elif isinstance(kv_cache, (list, tuple)) and all(
+                isinstance(t, torch.Tensor) for t in kv_cache
+            ):
+                for state_tensor in kv_cache:
+                    self._fill_state_tensor(state_tensor)
             else:
-                # Constant fill value
-                fill_values = torch.full(
-                    (len(valid_block_ids),) + block_shape,
-                    self.fill_mean,
-                    dtype=kv_cache.dtype,
-                    device=kv_cache.device,
+                logger.warning_once(
+                    "DecodeBenchConnector: skipping fill for layer %s whose KV "
+                    "cache is %s, not a tensor or a list/tuple of tensors.",
+                    layer_name,
+                    type(kv_cache).__name__,
                 )
-
-            # Batch fill operation
-            kv_cache[valid_block_ids] = fill_values
+                continue
 
         logger.debug(
             "DecodeBenchConnector: Filled %d blocks in group %d with %s values "
@@ -431,3 +419,62 @@ class DecodeBenchConnectorWorker:
             self.fill_mean,
             self.fill_std,
         )
+
+    def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
+        """Fill the requested block rows of a block-indexed KV cache tensor.
+
+        Args:
+            kv_cache: A KV cache tensor whose first dim is num_blocks.
+            block_ids: Block IDs to fill. IDs that are out of range for this
+                tensor's first dim are ignored.
+        """
+        # Convert block_ids to tensor on device
+        block_ids_tensor = torch.tensor(
+            block_ids, dtype=torch.long, device=kv_cache.device
+        )
+
+        # Filter invalid block IDs
+        valid_mask = block_ids_tensor < kv_cache.shape[0]
+        valid_block_ids = block_ids_tensor[valid_mask]
+
+        if len(valid_block_ids) == 0:
+            return
+
+        # Create fill values - either constant or random
+        block_shape = kv_cache.shape[1:]
+        if self.fill_std > 0:
+            # Random normal sampling
+            fill_values = torch.normal(
+                mean=self.fill_mean,
+                std=self.fill_std,
+                size=(len(valid_block_ids),) + block_shape,
+                dtype=kv_cache.dtype,
+                device=kv_cache.device,
+            )
+        else:
+            # Constant fill value
+            fill_values = torch.full(
+                (len(valid_block_ids),) + block_shape,
+                self.fill_mean,
+                dtype=kv_cache.dtype,
+                device=kv_cache.device,
+            )
+
+        # Batch fill operation
+        kv_cache[valid_block_ids] = fill_values
+
+    def _fill_state_tensor(self, kv_cache: torch.Tensor):
+        """Fill an entire non-block-indexed state tensor with dummy values.
+
+        Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
+        store their per-layer state as tensors with no num_blocks dimension,
+        so the whole tensor is filled with the same constant or random values
+        used for block fills, rather than selected block rows.
+
+        Args:
+            kv_cache: A state tensor to fill in its entirety.
+        """
+        if self.fill_std > 0:
+            kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
+        else:
+            kv_cache.fill_(self.fill_mean)