Bases: Module, AttentionLayerBase
Raw [index_k | index_z] rows of each request's open CSA region.
A region summary is a fixed-order online softmax over its eight rows, so
the rows are the whole state: the row that closes a region replays the
other seven from here, and a rejected draft is undone by the next step
overwriting its row.
One verify step stores up to k + 1 rows while the next step may still
replay up to seven earlier rows of the same region, so the ring needs
7 + k rows. It is sized to one main KV page so that it shares the
page size of the other KV cache groups.
Source code in vllm/models/step5/nvidia/csa_tail_ring.py
| class Step5CSATailRing(nn.Module, AttentionLayerBase):
"""Raw ``[index_k | index_z]`` rows of each request's open CSA region.
A region summary is a fixed-order online softmax over its eight rows, so
the rows are the whole state: the row that closes a region replays the
other seven from here, and a rejected draft is undone by the next step
overwriting its row.
One verify step stores up to ``k + 1`` rows while the next step may still
replay up to seven earlier rows of the same region, so the ring needs
``7 + k`` rows. It is sized to one main KV page so that it shares the
page size of the other KV cache groups.
"""
def __init__(
self,
*,
prefix: str,
main_layer_name: str,
num_speculative_tokens: int,
static_forward_context: dict[str, Any],
):
super().__init__()
if prefix in static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
static_forward_context[prefix] = self
self._static_forward_context = static_forward_context
self.layer_name = prefix
self.main_layer_name = main_layer_name
self.required_rows = 7 + num_speculative_tokens
self.kv_cache = torch.tensor([])
self.rows = self.kv_cache
self.capacity = 0
self.rows_per_block = 0
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
main = self._static_forward_context[self.main_layer_name]
main_spec = main.get_kv_cache_spec(vllm_config)
page_bytes = main_spec.page_size_bytes
if page_bytes % _ROW_BYTES or page_bytes // _ROW_BYTES < self.required_rows:
raise ValueError(
f"Step5 CSA tail ring needs {self.required_rows} rows of "
f"{_ROW_BYTES} bytes in one {page_bytes}-byte KV page; "
"increase --block-size."
)
self.capacity = page_bytes // _ROW_BYTES
# Declare the main block size so the ring does not coarsen the hybrid
# prefix-cache alignment; the page holds ``capacity`` 1 KB rows.
block_size = main_spec.block_size
return CircularBufferSpec(
block_size=block_size,
num_kv_heads=1,
head_size=page_bytes // (block_size * torch.bfloat16.itemsize),
head_size_v=0,
dtype=torch.bfloat16,
# Rows are addressed across blocks, which may be strided.
block_stride_alignment=_ROW_BYTES,
)
def get_attn_backend(self) -> type[AttentionBackend]:
return Step5CSATailRingBackend
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
assert kv_cache.stride(0) % _ROW_ELEMENTS == 0
self.kv_cache = kv_cache
self.capacity = kv_cache[0].numel() // _ROW_ELEMENTS
self.rows_per_block = kv_cache.stride(0) // _ROW_ELEMENTS
# Block b's ring is rows [b * rows_per_block, + capacity).
self.rows = kv_cache.as_strided(
(
(kv_cache.shape[0] - 1) * self.rows_per_block + self.capacity,
_ROW_ELEMENTS,
),
(_ROW_ELEMENTS, 1),
)
def metadata(self) -> Step5CSATailRingMetadata:
attn_metadata = get_forward_context().attn_metadata
assert isinstance(attn_metadata, dict)
return cast(Step5CSATailRingMetadata, attn_metadata[self.layer_name])
|