vllm - 💡(How to fix) Fix [RFC]: Standardize KV-cache Layouts [2 pull requests]

[num_layers, num_blocks, num_states, num_heads, <state_content>]

---

@dataclass
class KVCacheSpec:
    num_heads: int          # heads (1 if headless)
    tokens_per_state: int   # -1 infinite, 1 standard, N compressed
    state_content_size: int # bytes per state per head

---

Semantic:  [L, B, S, H, C]
Physical:  [L, B, S, H, C]  (identity)
stride_order = (0, 1, 2, 3, 4)

---

Semantic:  [L, B, S, H, C]
Physical:  [L, B, H, S, C]
stride_order = (0, 1, 3, 2, 4)

---

Semantic:  [L, B, S, H, C]
Physical:  [B, L, S, H, C]
stride_order = (1, 0, 2, 3, 4)

---

Semantic:  [L, B, S, H, C]
Physical:  [B, H, L, S, C]
stride_order = (1, 3, 0, 2, 4)

---

class KVCacheSpec:
    def slice_for_tp_transfer(self, my_tensor, my_tp, my_rank,
                               other_spec, other_tp,
                               other_rank) -> list[Tensor]:
        """Return slices of `my_tensor` that need to be transferred.
        Takes both sides' TP configs since tp_size is a deployment
        property, not part of the spec."""
        ...

---

def slice_for_tp_transfer(self, src_tensor, src_tp, src_rank,
                          dst_spec, dst_tp, dst_rank):

    src_heads = src_tp*self.num_heads
    dst_heads = dst_tp*dst_spec.num_heads

    if dst_heads > src_heads:
        dst_rank %= src_tp
    elif src_rank >= dst_tp:
        return []

    head_range = lambda tp_rank, num_heads: (tp_rank * num_heads, (tp_rank + 1) * num_heads)

    src_head_range = head_range(src_rank, self.num_heads)
    dst_head_range = head_range(dst_rank, dst_spec.num_heads)

    overlap = get_overlap(src_head_range, dst_head_range)
    if overlap is None:
        return []

    # Convert global overlap to local tensor indices
    local_slice = slice(overlap.start - src_head_range[0],
                        overlap.stop - src_head_range[0])
    return [src_tensor[:, :, :, local_slice]]

def get_overlap(range1, range2):
    # Calculate the intersection boundaries
    overlap_start = max(range1[0], range2[0])
    overlap_end = min(range1[1], range2[1])
    
    # Check if the boundaries form a valid range
    if overlap_start <= overlap_end:
        return slice(overlap_start, overlap_end)
    return None

---

def slice_for_tp_transfer(self, my_tensor, my_tp, my_rank,
                          other_spec, other_tp, other_rank):
    # tensor: [layers, blocks, 1, 1, conv_dim/tp, kernel-1]
    x_range = self._compute_x_range(my_tp, my_rank, other_tp, other_rank)
    b_range = self._compute_b_range(my_tp, my_rank, other_tp, other_rank)
    c_range = self._compute_c_range(my_tp, my_rank, other_tp, other_rank)
    return [
        my_tensor[:, :, 0, 0, x_range, :],  # x slice
        my_tensor[:, :, 0, 0, b_range, :],  # B slice
        my_tensor[:, :, 0, 0, c_range, :],  # C slice
    ]

---

# TP slicing (once at setup — determines which heads to transfer)
src_slices = src_spec.slice_for_tp_transfer(
    src_cache, src_tp, src_rank, dst_spec, dst_tp, dst_rank)
dst_meta = make_meta_kv_cache(dst_spec, num_blocks, dst_block_size)
dst_slices = dst_spec.slice_for_tp_transfer(
    dst_meta, dst_tp, dst_rank, src_spec, src_tp, src_rank)

# Per-block transfer (handles hetero block sizes naturally)
for src_states_slice, dst_states_slice in block_mapping:
    for src_s, dst_s in zip(src_slices, dst_slices):
        src_block = src_s[src_states_slice]
        dst_block = dst_s[dst_states_slice]
        for src_chunk, dst_chunk in jointly_contiguous_chunks(
                src_block, dst_block):
            self._rdma_transfer(src_chunk, dst_chunk)