In-engine H2O attention-weight extraction + FlashInfer probes (stacked on #133)

[mmjerge]

Summary

Follow-up to awslabs#133 (awslabs vLLM V1 KV-cache bridge). This is the in-engine H2O attention-weight work @mseeger asked to review directly. It is stacked on vllm-integration (the awslabs#133 branch), so the diff here is only the incremental H2O-wiring code — not a re-review of awslabs#133.

What's here

• examples/vllm_h2o_probe.py — selects the FlashInfer backend via the vLLM 0.23 attention_backend knob and installs forward hooks on the 24 Attention layers (confirms hooks fire and output stays correct).
• examples/vllm_h2o_lse_probe.py — wraps FlashInferImpl.forward and maps the impl surface (Q/K/V, paged kv_cache, FlashInferMetadata fields, slot_mapping, can_return_lse_for_decode) — i.e. exactly what's reachable for scoring.
• examples/vllm_h2o_score_probe.py — the key one: for single-sequence decode, gathers the request's K/V from the paged cache and computes correct per-KV-position attention mass in-engine using reference_summed_attention (the task-3 contract from Add vLLM (V1) KV cache policy integration: lastrec bridge + spike awslabs/keys_values#133). Sanity check: per-KV-head mass ≈ query-heads-per-group, seq_len grows by 1 each step.
• docs/vllm_integration.md — new H2O wiring (task 4.3) section: Blocker 1 (weights in-engine, tractable — single FlashInfer call, LSE available via return_lse=True) and Blocker 2 (arbitrary middle-block eviction vs. vLLM's append-only block table — the genuinely hard part), plus the mapped FlashInfer impl surface and a recommended order/effort estimate.

How this connects to the weights question

The math for summed attention weights is already implemented and unit-tested in awslabs#133 (keys_values/vllm/attention.py, test/vllm/test_attn_weights.py) and returns weights in a single SDPA call. This PR is the in-engine wiring of that signal: probing where to hook, and a working per-position score computation for the single-seq decode case.

Tested

• Probes run on a GPU box (g5.xlarge / A10G, vLLM 0.23, Qwen2.5-0.5B): FlashInfer selectable, hooks fire, output correct; score probe produces a correct in-engine score signal for single-sequence decode. They require CUDA + vLLM (guarded with clear exits otherwise).
• black --check clean across examples/, keys_values/vllm, test/vllm.

Not included (follow-ups)

• Multi-request batching: correlating each of the T batched tokens back to its request/layer every step.
• Score → block mapping and the arbitrary middle-block eviction ("holes") problem (Blocker 2).

AI usage

See ai_dev/vllm_integration.md.