feat(connector): support vLLM hybrid memory allocator (SupportsHMA)

[moreWax]

What

Adds vLLM hybrid memory allocator (HMA) support to PegaKVConnector: models that mix sliding-window and full attention (Gemma-style 5:1 layouts and similar) currently cannot use PegaFlow at all — vLLM refuses to combine a non-SupportsHMA connector with the hybrid KV cache manager, and disabling the manager makes SWA models allocate full-window KV.

Closes #228. Builds on and supersedes #202 — thanks @xiaguan for the groundwork; this rework exists because the on-GPU group semantics turned out to be different from what #202 assumed (details below).

Design

Observed semantics on vLLM 0.22.1 (KVCacheBlocks, SlidingWindowManager, and the in-tree HMA reference SimpleCPUOffloadConnector):

• Per-group block lists are positional, not suffix-shortened. SWA groups keep full-length lists with the null block filling out-of-window positions. So all alignment here is by block position, shared with the (single) hash stream. fix(connector): align hybrid save window across kv groups #202's common-suffix arithmetic was replaced accordingly.
• Group counts are model-derived. A 48-layer 5:1 hybrid yields 6 groups (5 SWA + 1 full-attention), and the full-attention group is last — nothing assumes 2 groups or that group 0 is full attention. The locally-computed prefix is measured on the FA group (SWA groups under-count: masked positions lose their hashes).
• Loads are one RPC per kv_cache_group, each group's layers paired with that group's destination block ids. The engine consumes a query lease on first use, so the scheduler mints one lease per group (extra query_prefetch over the already-resident range); a failed mint degrades to reporting that group's real destinations as load errors, which vLLM recomputes.
• Null-padded destinations preserve the lease contract. vLLM trims full-prompt hits by one block (logits) and SWA groups don't allocate out-of-window positions; destination lists are padded with the null block id up to the lease length. The null block is never read by the model, so the extra DMA is harmless and len(lease hashes) == len(destinations) holds.
• Saves are gated per position on ALL groups via bind_gpu_block_pool, with three-state classification: valid positions are ref-pinned until finished_sending (the save copy is asynchronous — without the pin, a sliding-window block can slide out and be reused mid-DMA, storing foreign bytes under a valid hash); transient positions (allocated but not yet sealed/hashed) stop the cursor and retry next round instead of punching permanent holes in the hash chain; dead positions (slid out / reused — pool hash carries vLLM's 4-byte group-id suffix, stripped and verified) are skipped for all groups together so every stored hash is complete across all layers.
• Partial group failures surface exactly once. Loads being per-group RPCs means one group can fail while another is in flight; failed groups' blocks are recompute-marked immediately, but the request is reported in finished_recving only after its in-flight group states drain (a premature duplicate report trips vLLM's scheduler asserts). Fully-out-of-window SWA slices (all-null destinations) skip their RPC and release the minted lease.
• Non-attention groups are refused loudly. In multi-group configs only FullAttentionSpec + SlidingWindowSpec are accepted — Mamba/SSM or cross-attention state must not be saved/loaded as positional KV. Single-group models are untouched (including MLA page-first striping, which stays on the legacy path).
• Safety rails: the storage namespace gains a kv-group-signature factor (two group layouts must never share blocks); cross-layer registration is forced off for multi-group; heterogeneous per-group block sizes are rejected loudly at init — supporting them needs per-group hash granularity and is called out as future work rather than silently corrupting hash↔block mappings.
• request_finished_all_groups is implemented; world_size stays the true worker count (after fix(core)!: seal instance layer topology from registered layers #346 the server seals topology from registered layers, so fix(connector): align hybrid save window across kv groups #202's world_size * num_groups is no longer needed).

Validation

• Full python unit suite passes; new tests cover positional multi-group saves, shortest-group capping, external-hit positional indexing, a 6-group SWA-hybrid save with a reused block (group-id-suffixed pool hashes, permanent skip), transient-position cursor stop/resume, ref-pin bookkeeping across finished_sending, exactly-once reporting under partial group failure, all-null-slice lease release, and partial-accept destination padding.
• Live on vLLM 0.22.1 (RTX 3090):
  • dense model: cross-restart prefix reuse, 335/335 blocks, ~0.2 s vs ~2.5 s cold;
  • 48-layer 5:1 SWA hybrid (hybrid manager disabled, single-group — the heterogeneous-block-size case): cross-restart reuse 264/265 blocks, 0.52 s vs 4.09 s cold, byte-identical output.
• We were not able to run the repo's exact test_vllm_e2e_correctness.py gate (no /data/models/Qwen3-4B here); happy to iterate if CI or a reviewer run surfaces anything.

Limitations / future work

• Heterogeneous per-group block sizes (e.g. a full-attention group with fewer bytes/token getting larger blocks under uniform-page-size grouping) are detected and refused. Supporting them needs per-group hash striding and likely per-group-class layer topologies.
• Out-of-window SWA loads currently DMA into the null block (correct but wasted bandwidth on long prefixes); trailing-window leases per SWA group would trim that.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR adds vLLM Hybrid Memory Allocator (HMA) compatibility to the PegaFlow vLLM connector by making load/save intent metadata group-aware (per kv_cache_group) and updating both scheduler- and worker-side logic to correctly handle multi-group hybrid layouts (e.g., sliding-window + full attention).

Changes:

• Implement HMA support in PegaKVConnector via SupportsHMA, including per-group load RPCs and multi-group save semantics.
• Extend connector metadata (LoadIntent, SaveIntent, PegaConnectorMetadata) to carry per-group block ids and leases plus layer→group mapping.
• Add/extend tests covering multi-group behaviors like partial group failures, null-slice lease release, and positional save alignment across groups.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
python/pegaflow/connector/init.py	Declares SupportsHMA, wires kv-cache-group config into scheduler connector, and adds group signature + block size validation for safe HMA operation.
python/pegaflow/connector/common.py	Updates intent metadata structures to be group-aware; adds namespace isolation factor for group layout signature.
python/pegaflow/connector/scheduler.py	Implements scheduler-side HMA behavior (group parsing, per-group leases, positional multi-group save cutting with pool validation and pinning).
python/pegaflow/connector/worker.py	Implements worker-side per-group load tracking and HMA-aware save batching per layer/group.
python/tests/test_combine_hashes.py	Updates existing tests to the new group_block_ids shape and adds extensive hybrid/HMA behavioral coverage.
python/tests/test_connector_fault_tolerance.py	Adds tests for HMA-specific fault tolerance (partial group failures and null-slice lease release).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[moreWax]

Fixed in 8ff0b0e — KeyError added to the except tuple so dict-like pools classify missing ids as dead instead of aborting the intent cut.

[moreWax]

Fixed in 8ff0b0e — the pool-unbound branch no longer consumes the cursor; the range is retried once the pool binds.

[moreWax]

Fixed in 8ff0b0e — an out-of-range group now warns and skips the layer; no more silent remap to group 0.

[moreWax]

Fixed in 8ff0b0e — dead states_by_req removed.

[moreWax]

Addressed all review feedback in 8ff0b0e: dict-like pool misses (KeyError) classify as dead; the save cursor is kept (not consumed) when the block pool is not yet bound; an out-of-range layer group now warns and skips instead of silently remapping to group 0; dead states_by_req removed. Full unit suite green.

[xiaguan]

Thanks for the detailed writeup — the positional multi-group semantics analysis (null-block padding, FA-group prefix accounting, packed BlockHashWithGroupId) is genuinely useful groundwork.

Before we go deeper on review, could you share the concrete motivation behind this?

1. Which hybrid-attention model(s) do you actually need PegaFlow for, and in what deployment? Is there a real workload waiting on this, or is it primarily completeness for feat: pegaflow vLLM connector supports HMA #228?
2. Do you have measurements on the true multi-group path? The live validation ran the 5:1 hybrid with the hybrid manager disabled (single-group) — which is exactly the configuration this PR's headline machinery doesn't cover. We'd want prefix-hit rate and TTFT-vs-cold numbers for a real multi-group SWA hybrid before weighing the trade-offs.
3. What does the end-to-end win look like on long prefixes? SWA groups DMA the full leased range with out-of-window positions landing in the null block, so a 5:1 layout spends roughly 5/6 of load bandwidth on bytes the model never reads. TTFT reduction is the whole point of loading external KV, so this matters a lot for whether the feature pays off.

Context for the questions: we deliberately kept the connector single-group so far because HMA support carries real ongoing cost — per-group lease mints as sequential blocking RPCs in the scheduler step, dependence on vLLM block-pool internals (packed hash layout, ref-pin APIs) that drift across versions, and ref-pinning GPU blocks across async saves, which works against the memory savings SWA is supposed to provide. Whether that cost is worth carrying long-term depends on a concrete production need rather than API completeness.

Also note: per repo policy, connector-visible changes need the test_vllm_e2e_correctness.py gate run before merge — we can rerun it on our side, but only once the motivation/benefit question is settled.