feat: add staged page-level prefix reuse

[Andrewxu313]

Add staged page-level prefix reuse

Goal

Implement opt-in, page-level prefix KV reuse for BatchGen. Do not implement token-level radix granularity or partial-page sharing in this PR.

The staged outcome is:

1. Host KV-cache page reuse for host page efficiency.
2. Prefix-aware prefill that computes and offloads only non-hit suffix pages/tokens.
3. Decode-batch GPU page materialization that loads each needed shared host page once and points multiple sequence page-table rows at the same physical GPU page where safe.

Background

The older prefix-cache approach must not be revived directly. It is prior-art risk context only.

This work should build on the request-pool and dynamic host-KV infrastructure, including request-pool scheduling, host-KV reservation/growth, host-KV eviction/re-entry, GPU on-hold/resume, and repeated-prefix workloads.

BatchGen-specific findings:

• For BatchGen's first stage, implement page-level reuse only. Use token IDs only to hash/validate full pages; do not split/share inside a page and do not share partial pages.
• BatchGen host KV is currently single-owner: host_paged_kv_backend.cpp tracks page_owners[page] = sequence_id, and ReleaseSequence() frees all pages for that sequence.
• BatchGen host page table is currently a flat sequence_id -> SequenceRecord{pages} and cannot represent [shared prefix pages + private suffix pages].
• BatchGen GPU KV is currently single-owner: gpu_paged_kv_manager.py stores one linear page vector per sequence, and release returns all pages for that sequence.
• Prefill currently recomputes and offloads full prompts. GPT-OSS/GQA is the first target because its prefill/decode page-table interfaces are clearest.

Design

Implement this as a staged, page-level feature. Each milestone must be gated so reviewers can distinguish host memory savings, prefill compute/offload savings, and GPU page-pressure savings.

Milestone 1: Host KV-cache page reuse

Implement host-side page reuse and host page efficiency first:

• Add a page-level prefix index keyed by full-page chained hashes: (model/cache namespace, page_size, page_index, parent_page_hash, current_page_token_hash).
• Add host page refcounts and prefix-entry pins.
• Extend host page-table records to represent shared prefix pages, private suffix pages, combined logical page view, shared_prefix_tokens, and private start position.
• Add a prefix-aware host allocation/binding API such as allocate_pages_for_sequences_with_prefix(requests).
• In prefill admission, lookup page-level prefix hits; for hits, attach shared host pages and allocate only suffix host pages.
• Make ReleaseSequence(), host unregister/release, allocation rollback, AsyncOffloadLayerKVToHost(), and AsyncAppendDecodeKVToHost() shared-page safe.
• Add host stats for logical pages, physical pages, shared pages, private pages, refcount increments/decrements, and host pages saved.

Milestone 2: Prefill compute/offload reduction

After Milestone 1 host rows are correct, reduce prefill work for prefix hits:

• Build suffix-only prefill metadata: prefix_shared_tokens, suffix_input_ids, suffix_start_pos, suffix_length, and full logical context length.
• Preserve absolute position IDs/RoPE offsets.
• For GPT-OSS/GQA, compute Q/K/V only for suffix tokens.
• Make suffix attention attend over cached prefix KV pages plus newly computed suffix K/V.
• If the current FlashAttention path cannot consume paged prefix KV plus suffix K/V directly, gather matched prefix pages into a temporary batch-local KV view for prefill only.
• Offload only suffix K/V into private host pages; do not copy cache-hit prefix pages.
• Define exact full-hit behavior for suffix_length == 0 with a decode-like or cached-prefill path. Do not silently fall back to full prefill.

Milestone 3: Decode-batch GPU page materialization

After the decode batch is formed, reduce GPU page pressure:

• Build a decode-batch plan from each sequence's logical host row [shared prefix host pages + private suffix/decode host pages].
• Deduplicate host pages within the rank/decode batch.
• Allocate/load one GPU page per unique missing host page.
• Extend GPU sequence state and page-table rebuild so multiple sequence rows can reference the same physical GPU prefix page.
• Keep decode runway/new-token pages private.
• Make GPU release refcount-aware for IN_DECODE -> ON_HOLD, ON_HOLD -> IN_DECODE, completion, and extension failure.
• Add stats for unique host pages loaded, skipped duplicate H2D pages, logical GPU pages, physical GPU pages, and shared GPU refcounts.

Acceptance criteria

• Feature is disabled by default behind --enable-prefix-reuse.
• Disabled-feature behavior preserves existing request-pool and dynamic-host-KV behavior.
• Prefix matching is page-level only; no token-level split, no partial-page sharing, and no raw prompt text keying.
• Milestone 1: host KV supports shared prefix pages with correct page refcounts, prefix-entry pins, allocation rollback, and release behavior.
• Milestone 1: host page-table rows can represent [shared prefix pages + private suffix pages] while preserving the legacy combined Pages(sequence_id) view.
• Milestone 1: host stats show fewer physical host pages than logical pages for repeated full-page prefixes.
• Milestone 2: GPT-OSS/GQA prefix-hit prefill computes and offloads only suffix pages/tokens while preserving correctness versus full prefill within accepted tolerance.
• Milestone 2: exact full-hit behavior is implemented or explicitly guarded with a clear error; no silent full-prefill fallback.
• Milestone 3: decode-batch planning deduplicates shared host pages and loads each unique needed page to GPU once per rank/batch.
• Milestone 3: GPU page-table rows can reference shared physical prefix pages and private decode pages safely.
• Decode writes never target shared prefix pages.
• Lifecycle transitions are refcount-safe for PREFILLED, IN_DECODE, ON_HOLD, EVICTED, and COMPLETED.
• Prefix-hit and prefix-miss sequences can coexist in the same prefill and decode batches.
• DSA/MLA unsupported paths fail loudly or are explicitly gated unless this task is expanded.
• Tests cover page-level lookup, host refcounts, GPU refcounts, suffix-only prefill correctness, decode-batch GPU sharing, IN_DECODE -> ON_HOLD -> IN_DECODE, host eviction/re-entry, and completion release.
• GPU validation passes on an approved GPU host with mandatory clean/verify before server launch.