feat: add staged page-level prefix reuse

batchgen/attention/forward_metadata.py

@@ -0,0 +1,152 @@
+"""First-class forward metadata for attention execution.
+
+These dataclasses describe the logical forward batch without depending on
+legacy wrapper class variables. They intentionally do not mutate runtime state;
+builders are responsible for constructing them from already-validated static
+planning inputs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+import torch
+
+ForwardPhase = Literal["prefill", "decode"]
+
+
+@dataclass(frozen=True)
+class PrefillAttentionMetadata:
+    """Attention metadata for prefill or suffix-only prefill.
+
+    Prefix reuse is represented by q/kv length divergence:
+    ``kv_seq_lens[i] - q_seq_lens[i]`` is the cached prefix length for sequence
+    ``i``. The legacy wrapper context mirrors these derived values into
+    ``AttnWrapperBase.prepack_prefix_*`` for model wrappers.
+    """
+
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+    max_seqlen_q: int
+    max_seqlen_k: int
+    q_seq_lens: list[int]
+    kv_seq_lens: list[int]
+    position_ids: torch.Tensor
+    append_seq_lens: Optional[list[int]] = None
+
+    @property
+    def batch_size(self) -> int:
+        return len(self.q_seq_lens)
+
+
+@dataclass(frozen=True)
+class DecodeAttentionMetadata:
+    """Attention metadata for decode forward batches."""
+
+    cache_seqlens: torch.Tensor
+    max_seqlen: int
+    page_table: Optional[torch.Tensor] = None
+    slot_indices: Optional[torch.Tensor] = None
+    batch_slice: Optional[slice] = None
+
+    @property
+    def batch_size(self) -> int:
+        return int(self.cache_seqlens.numel())
+
+
+@dataclass(frozen=True)
+class KVCacheMetadata:
+    """KV cache handles associated with a forward batch."""
+
+    gpu_paged_kv_manager: Optional[object] = None
+    host_worker_view: Optional[object] = None
+    aux_gpu_paged_kv_manager: Optional[object] = None
+    aux_host_worker_view: Optional[object] = None
+    prefill_prefix_materialization: Optional[object] = None
+
+
+@dataclass(frozen=True)
+class ForwardBatchMetadata:
+    """Top-level metadata object for one model forward batch."""
+
+    phase: ForwardPhase
+    global_sequence_ids: list[int]
+    prefill: Optional[PrefillAttentionMetadata] = None
+    decode: Optional[DecodeAttentionMetadata] = None
+    kv_cache: Optional[KVCacheMetadata] = None
+
+    def require_prefill(self) -> PrefillAttentionMetadata:
+        if self.phase != "prefill" or self.prefill is None:
+            raise RuntimeError(
+                "Prefix cache prepack metadata requires prefill metadata"
+            )
+        return self.prefill
+
+    @property
+    def cu_seqlens(self) -> torch.Tensor:
+        return self.require_prefill().cu_seqlens_q
+
+    @property
+    def cu_seqlens_cpu(self) -> list[int]:
+        return _build_cu_seqlens_values(self.seq_lengths)
+
+    @property
+    def max_seqlen(self) -> int:
+        return int(self.require_prefill().max_seqlen_q)
+
+    @property
+    def num_sequences(self) -> int:
+        return int(self.require_prefill().batch_size)
+
+    @property
+    def seq_lengths(self) -> list[int]:
+        return [int(length) for length in self.require_prefill().q_seq_lens]
+
+    @property
+    def append_seq_lengths(self) -> list[int]:
+        prefill = self.require_prefill()
+        if prefill.append_seq_lens is None:
+            return [int(length) for length in prefill.q_seq_lens]
+        return [int(length) for length in prefill.append_seq_lens]
+
+    @property
+    def prefix_shared_tokens(self) -> Optional[list[int]]:
+        tokens = [
+            int(kv_len) - int(append_len)
+            for kv_len, append_len in zip(
+                self.require_prefill().kv_seq_lens,
+                self.append_seq_lengths,
+            )
+        ]
+        if any(token < 0 for token in tokens):
+            raise RuntimeError(
+                "Prefix cache metadata requires kv lengths >= append lengths"
+            )
+        return tokens if any(token > 0 for token in tokens) else None
+
+    @property
+    def prefix_reuse_mode(self) -> bool:
+        tokens = self.prefix_shared_tokens
+        return tokens is not None and any(token > 0 for token in tokens)
+
+    @property
+    def full_seq_lengths(self) -> Optional[list[int]]:
+        if not self.prefix_reuse_mode:
+            return None
+        return [int(length) for length in self.require_prefill().kv_seq_lens]
+
+    def cu_seqlens_list(self) -> list[int]:
+        return list(self.cu_seqlens_cpu)
+
+    def append_seq_lengths_list(self) -> list[int]:
+        return list(self.append_seq_lengths)
+
+
+def _build_cu_seqlens_values(seq_lengths: list[int]) -> list[int]:
+    values = [0]
+    running = 0
+    for length in seq_lengths:
+        running += int(length)
+        values.append(running)
+    return values

batchgen/attention/forward_metadata_context.py

@@ -0,0 +1,177 @@
+"""Context binding for first-class attention forward metadata.
+
+This module is the compatibility bridge between explicit
+``ForwardBatchMetadata`` and the legacy ``AttnWrapperBase`` class variables.
+The metadata object remains the source of truth; legacy fields are only
+populated for the dynamic extent of a single forward call.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Iterator, Optional
+
+from batchgen.attention.forward_metadata import (
+    DecodeAttentionMetadata,
+    ForwardBatchMetadata,
+    KVCacheMetadata,
+    PrefillAttentionMetadata,
+)
+
+
+_CURRENT_FORWARD_BATCH_METADATA: ContextVar[Optional[ForwardBatchMetadata]] = (
+    ContextVar("current_forward_batch_metadata", default=None)
+)
+
+_LEGACY_ATTENTION_FIELDS = (
+    "phase",
+    "cur_batch",
+    "position_ids",
+    "prepack_mode",
+    "prepack_cu_seqlens",
+    "prepack_max_seqlen",
+    "prepack_num_sequences",
+    "prepack_seq_lengths",
+    "prepack_append_seq_lengths",
+    "prepack_prefix_reuse_mode",
+    "prepack_prefix_shared_tokens",
+    "prepack_full_seq_lengths",
+    "cache_seqlens",
+    "max_seqlen",
+    "gpu_paged_kv_manager",
+    "host_paged_kv_worker_view",
+    "prefill_prefix_materialization",
+    "gpu_paged_kv_manager_aux",
+    "host_paged_kv_worker_view_aux",
+)
+
+
+def get_current_forward_batch_metadata(
+    required: bool = False,
+) -> Optional[ForwardBatchMetadata]:
+    """Return the metadata bound to the current execution context."""
+
+    metadata = _CURRENT_FORWARD_BATCH_METADATA.get()
+    if metadata is None and required:
+        raise RuntimeError("ForwardBatchMetadata is required but is not bound")
+    return metadata
+
+
+@contextmanager
+def bind_forward_batch_metadata(
+    metadata: ForwardBatchMetadata,
+) -> Iterator[ForwardBatchMetadata]:
+    """Bind metadata for one forward and mirror it into legacy wrapper fields."""
+
+    if not isinstance(metadata, ForwardBatchMetadata):
+        raise TypeError("metadata must be a ForwardBatchMetadata instance")
+
+    # Import lazily so metadata users can be unit-tested without importing model
+    # wrappers unless the compatibility bridge is actually used.
+    from batchgen.models.wrappers.attention import AttnWrapperBase
+
+    previous_values = {
+        field: getattr(AttnWrapperBase, field, None)
+        for field in _LEGACY_ATTENTION_FIELDS
+    }
+    token = _CURRENT_FORWARD_BATCH_METADATA.set(metadata)
+    try:
+        _sync_legacy_attention_wrapper(AttnWrapperBase, metadata)
+        yield metadata
+    finally:
+        _CURRENT_FORWARD_BATCH_METADATA.reset(token)
+        for field, value in previous_values.items():
+            setattr(AttnWrapperBase, field, value)
+
+
+def _sync_legacy_attention_wrapper(
+    wrapper_cls: type,
+    metadata: ForwardBatchMetadata,
+) -> None:
+    wrapper_cls.phase = metadata.phase
+    wrapper_cls.cur_batch = list(metadata.global_sequence_ids)
+
+    if metadata.phase == "prefill":
+        assert metadata.prefill is not None
+        _sync_prefill_fields(wrapper_cls, metadata.prefill)
+    else:
+        assert metadata.decode is not None
+        _sync_decode_fields(wrapper_cls, metadata.decode)
+
+    if metadata.kv_cache is not None:
+        _sync_kv_cache_fields(wrapper_cls, metadata.kv_cache)
+
+
+def _sync_prefill_fields(
+    wrapper_cls: type,
+    prefill: PrefillAttentionMetadata,
+) -> None:
+    wrapper_cls.position_ids = prefill.position_ids
+    wrapper_cls.prepack_mode = True
+    wrapper_cls.prepack_cu_seqlens = prefill.cu_seqlens_q
+    wrapper_cls.prepack_max_seqlen = int(prefill.max_seqlen_q)
+    wrapper_cls.prepack_num_sequences = prefill.batch_size
+    wrapper_cls.prepack_seq_lengths = list(prefill.q_seq_lens)
+    wrapper_cls.prepack_append_seq_lengths = _append_seq_lens(prefill)
+    wrapper_cls.cache_seqlens = None
+    wrapper_cls.max_seqlen = None
+
+    _sync_prefix_reuse_fields(wrapper_cls, prefill)
+
+
+def _sync_prefix_reuse_fields(
+    wrapper_cls: type,
+    prefill: PrefillAttentionMetadata,
+) -> None:
+    prefix_lens = [
+        int(kv_len) - int(q_len)
+        for q_len, kv_len in zip(prefill.q_seq_lens, prefill.kv_seq_lens)
+    ]
+    if any(length < 0 for length in prefix_lens):
+        raise ValueError(
+            "prefill kv sequence lengths must be >= query sequence lengths"
+        )
+    if not any(length > 0 for length in prefix_lens):
+        wrapper_cls.prepack_prefix_reuse_mode = False
+        wrapper_cls.prepack_prefix_shared_tokens = None
+        wrapper_cls.prepack_full_seq_lengths = None
+        return
+
+    full_seq_lens = [int(length) for length in prefill.kv_seq_lens]
+    wrapper_cls.prepack_prefix_reuse_mode = True
+    wrapper_cls.prepack_prefix_shared_tokens = prefix_lens
+    wrapper_cls.prepack_full_seq_lengths = full_seq_lens
+
+
+def _sync_decode_fields(
+    wrapper_cls: type, decode: DecodeAttentionMetadata
+) -> None:
+    wrapper_cls.position_ids = None
+    wrapper_cls.prepack_mode = False
+    wrapper_cls.prepack_cu_seqlens = None
+    wrapper_cls.prepack_max_seqlen = None
+    wrapper_cls.prepack_num_sequences = None
+    wrapper_cls.prepack_seq_lengths = None
+    wrapper_cls.prepack_append_seq_lengths = None
+    wrapper_cls.prepack_prefix_reuse_mode = False
+    wrapper_cls.prepack_prefix_shared_tokens = None
+    wrapper_cls.prepack_full_seq_lengths = None
+    wrapper_cls.cache_seqlens = decode.cache_seqlens
+    wrapper_cls.max_seqlen = int(decode.max_seqlen)
+
+
+def _sync_kv_cache_fields(wrapper_cls: type, kv_cache: KVCacheMetadata) -> None:
+    wrapper_cls.gpu_paged_kv_manager = kv_cache.gpu_paged_kv_manager
+    wrapper_cls.host_paged_kv_worker_view = kv_cache.host_worker_view
+    wrapper_cls.prefill_prefix_materialization = (
+        kv_cache.prefill_prefix_materialization
+    )
+    wrapper_cls.gpu_paged_kv_manager_aux = kv_cache.aux_gpu_paged_kv_manager
+    wrapper_cls.host_paged_kv_worker_view_aux = kv_cache.aux_host_worker_view
+
+
+def _append_seq_lens(prefill: PrefillAttentionMetadata) -> list[int]:
+    if prefill.append_seq_lens is None:
+        return list(prefill.q_seq_lens)
+    return [int(length) for length in prefill.append_seq_lens]

batchgen/attention/gqa/__init__.py

@@ -6,6 +6,7 @@
 Key components:
 - gqa_prefill_fa: Prefill using flash_attn_varlen_func (unpadded sequences)
 - gqa_decode_fa: Decode using flash_attn_with_kvcache (paged KV cache)
+- gqa_extend_fa: Extend prefill using flash_attn_with_kvcache (paged KV cache)
 - apply_sink_correction: Post-correction for attention sinks
 - attention_ref: Reference implementation for testing
 
@@ -16,6 +17,7 @@
 
 from .fa_prefill import gqa_prefill_fa
 from .fa_decode import gqa_decode_fa, gqa_decode_fa_contiguous
+from .fa_extend import gqa_extend_fa
 from .sink_correction import apply_sink_correction
 from .reference import attention_ref, attention_ref_no_sinks
 from .gqa_mode3 import gqa_decoding_mode_3_bf16
@@ -25,6 +27,7 @@
     'gqa_prefill_fa',
     'gqa_decode_fa',
     'gqa_decode_fa_contiguous',
+    'gqa_extend_fa',
     'apply_sink_correction',
     'attention_ref',
     'attention_ref_no_sinks',