[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>
Dao007forever committed
Jun 22, 2026 at 11:54 UTC
3ce15fd574960faaf19fb202851bc140d0eb9d02
1 file changed
+79
-32
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+79
-32
@@ -387,40 +387,28 @@ class DecodeBenchConnectorWorker:
387
388
kv_cache = self.kv_caches[layer_name]
389
390
- # Convert block_ids to tensor on device
391
- block_ids_tensor = torch.tensor(
392
- block_ids, dtype=torch.long, device=kv_cache.device
393
- )
394
-
395
- # Filter invalid block IDs
396
- valid_mask = block_ids_tensor < kv_cache.shape[0]
397
- valid_block_ids = block_ids_tensor[valid_mask]
398
-
399
- if len(valid_block_ids) == 0:
400
- continue
401
-
402
- # Create fill values - either constant or random
403
- block_shape = kv_cache.shape[1:]
404
- if self.fill_std > 0:
405
- # Random normal sampling
406
- fill_values = torch.normal(
407
- mean=self.fill_mean,
408
- std=self.fill_std,
409
- size=(len(valid_block_ids),) + block_shape,
410
- dtype=kv_cache.dtype,
411
- device=kv_cache.device,
412
- )
390
+ # Attention layers store KV as a single block-indexed tensor whose
391
+ # first dim is num_blocks; fill the requested block rows. Hybrid /
392
+ # linear-attention layers (e.g. Mamba, Kimi Delta Attention) store
393
+ # their state as a list/tuple of tensors that are NOT block-indexed
394
+ # — each tensor is a single state buffer with no num_blocks
395
+ # dimension — so fill each tensor in its entirety with the same
396
+ # dummy values.
397
+ if isinstance(kv_cache, torch.Tensor):
398
+ self._fill_block_tensor(kv_cache, block_ids)
399
+ elif isinstance(kv_cache, (list, tuple)) and all(
400
+ isinstance(t, torch.Tensor) for t in kv_cache
401
+ ):
402
+ for state_tensor in kv_cache:
403
+ self._fill_state_tensor(state_tensor)
404
else:
414
- # Constant fill value
415
- fill_values = torch.full(
416
- (len(valid_block_ids),) + block_shape,
417
- self.fill_mean,
418
- dtype=kv_cache.dtype,
419
- device=kv_cache.device,
405
+ logger.warning_once(
406
+ "DecodeBenchConnector: skipping fill for layer %s whose KV "
407
+ "cache is %s, not a tensor or a list/tuple of tensors.",
408
+ layer_name,
409
+ type(kv_cache).__name__,
410
)
421
-
422
- # Batch fill operation
423
- kv_cache[valid_block_ids] = fill_values
411
+ continue
412
413
logger.debug(
414
"DecodeBenchConnector: Filled %d blocks in group %d with %s values "
@@ -431,3 +419,62 @@ class DecodeBenchConnectorWorker:
419
self.fill_mean,
420
self.fill_std,
421
)
422
+
423
+ def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
424
+ """Fill the requested block rows of a block-indexed KV cache tensor.
425
+
426
+ Args:
427
+ kv_cache: A KV cache tensor whose first dim is num_blocks.
428
+ block_ids: Block IDs to fill. IDs that are out of range for this
429
+ tensor's first dim are ignored.
430
+ """
431
+ # Convert block_ids to tensor on device
432
+ block_ids_tensor = torch.tensor(
433
+ block_ids, dtype=torch.long, device=kv_cache.device
434
+ )
435
+
436
+ # Filter invalid block IDs
437
+ valid_mask = block_ids_tensor < kv_cache.shape[0]
438
+ valid_block_ids = block_ids_tensor[valid_mask]
439
+
440
+ if len(valid_block_ids) == 0:
441
+ return
442
+
443
+ # Create fill values - either constant or random
444
+ block_shape = kv_cache.shape[1:]
445
+ if self.fill_std > 0:
446
+ # Random normal sampling
447
+ fill_values = torch.normal(
448
+ mean=self.fill_mean,
449
+ std=self.fill_std,
450
+ size=(len(valid_block_ids),) + block_shape,
451
+ dtype=kv_cache.dtype,
452
+ device=kv_cache.device,
453
+ )
454
+ else:
455
+ # Constant fill value
456
+ fill_values = torch.full(
457
+ (len(valid_block_ids),) + block_shape,
458
+ self.fill_mean,
459
+ dtype=kv_cache.dtype,
460
+ device=kv_cache.device,
461
+ )
462
+
463
+ # Batch fill operation
464
+ kv_cache[valid_block_ids] = fill_values
465
+
466
+ def _fill_state_tensor(self, kv_cache: torch.Tensor):
467
+ """Fill an entire non-block-indexed state tensor with dummy values.
468
+
469
+ Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
470
+ store their per-layer state as tensors with no num_blocks dimension,
471
+ so the whole tensor is filled with the same constant or random values
472
+ used for block fills, rather than selected block rows.
473
+
474
+ Args:
475
+ kv_cache: A state tensor to fill in its entirety.
476
+ """
477
+ if self.fill_std > 0:
478
+ kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
479
+ else:
480
+ kv_cache.fill_(self.fill_mean)