From 80cbe998f3b7b4311e2a25d52d42330a61216a11 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 14:36:34 -0400
Subject: [PATCH 1/8] Record task 4.3 wiring investigation: LSE weight path +
 non-prefix eviction options

---
 docs/vllm_integration.md | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/docs/vllm_integration.md b/docs/vllm_integration.md
index 912a0af..6de01fa 100644
--- a/docs/vllm_integration.md
+++ b/docs/vllm_integration.md
@@ -264,6 +264,42 @@ Open follow-up: this uses monkeypatches in an experiment harness. A production
 path should register via the `register_custom_kv_cache_specs` platform hook and a
 config flag rather than patching `get_kv_cache_spec`/layer attributes.
 
+## H2O wiring (task 4.3): findings and plan
+
+Two sub-problems must be solved to run H2O in-engine; investigation status:
+
+### Blocker 1 - attention weights in-engine (tractable)
+H2O scores need per-KV-position attention mass. FlashAttention (the default
+backend) exposes nothing. But vLLM's **V1 FlashInfer backend** runs its wrappers
+with `return_lse=True` (`vllm/v1/attention/backends/flashinfer.py`), yielding the
+per-query log-sum-exp. With `LSE + Q + K`, per-position weights are
+`exp(q_i·k_j·scale - lse_i)` - the same Q,K,LSE recipe keys_values' FlashInfer
+path already uses, and summable with the existing Triton score-sum. So the plan
+is: force `VLLM_ATTENTION_BACKEND=FLASHINFER`, capture LSE per layer, compute the
+per-position score-sum, and feed `H2OManager.record_block_scores`.
+
+### Blocker 2 - non-prefix eviction in the block table (the hard part)
+vLLM's built-in managers only ever free a contiguous **prefix** (sliding window,
+chunked-local), replacing it with null blocks; attention then runs over the
+suffix. H2O evicts arbitrary **middle** blocks, leaving holes the block table +
+FlashInfer masking don't natively represent, and V1 block tables are designed
+append-only. Candidate approaches, to evaluate empirically:
+- **A. Block-table compaction**: drop the evicted block from the request's list.
+  The stored K keeps its original RoPE (position is baked into K, not the list
+  index), so attention over the remaining blocks uses correct encodings - this
+  is promising, but `slot_mapping`/`seq_len`/causal-mask metadata must be made
+  consistent, against an append-only assumption.
+- **B. Dense-cache attention backend (Option B)**: faithful per-head eviction,
+  most invasive.
+
+### Recommended 4.3 order
+1. Wire blocker 1 first (LSE capture -> score-sum -> `record_block_scores`) and
+   verify scores look sane during generation; this is reusable regardless of how
+   eviction is wired.
+2. Prototype approach A (compaction) on short sequences; compare against the
+   keys_values/LitGPT H2O reference. If compaction can't be made consistent with
+   V1 metadata, escalate to Option B.
+
 ## Open questions
 
 - Can a `SingleTypeKVCacheManager` express keys_values eviction without touching

From 6f26cb9793880c5c43b7e548dec3af596e84360b Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 14:50:22 -0400
Subject: [PATCH 2/8] Add task 4.3 probe: FlashInfer backend + attention-layer
 hooks (step 1)

---
 examples/vllm_h2o_probe.py | 161 +++++++++++++++++++++++++++++++++++++
 1 file changed, 161 insertions(+)
 create mode 100644 examples/vllm_h2o_probe.py

diff --git a/examples/vllm_h2o_probe.py b/examples/vllm_h2o_probe.py
new file mode 100644
index 0000000..44a4888
--- /dev/null
+++ b/examples/vllm_h2o_probe.py
@@ -0,0 +1,161 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Task 4.3 probe, step 1: foundation for in-engine H2O scoring.
+
+Before the hard parts (capturing per-position attention weights via the
+FlashInfer backend's LSE, and reflecting evicted blocks in the attention path),
+this probe verifies the plumbing we depend on:
+
+1. vLLM can run with the FlashInfer backend on this GPU
+   (``VLLM_ATTENTION_BACKEND=FLASHINFER``), which is the backend that exposes
+   LSE (FlashAttention does not).
+2. We can install forward hooks on the model's ``Attention`` layers and they
+   fire during generation, so a later step can capture the tensors needed to
+   compute H2O scores.
+
+It logs, for the first few hook fires, the layer name and the query/output
+shapes. It does not yet compute scores. Runs the engine in-process so hooks
+apply in this interpreter.
+
+Usage:
+    python examples/vllm_h2o_probe.py --model Qwen/Qwen2.5-0.5B-Instruct
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+os.environ.setdefault("VLLM_ATTENTION_BACKEND", "FLASHINFER")
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+import sys
+
+_MAX_LOGS = 6
+_log_count = 0
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    p.add_argument("--max-model-len", type=int, default=2048)
+    p.add_argument(
+        "--prompt",
+        default="Count to five, then explain why selective KV caching helps.",
+    )
+    p.add_argument("--max-tokens", type=int, default=16)
+    return p.parse_args()
+
+
+def _check_env() -> None:
+    try:
+        import torch
+    except ImportError:
+        sys.exit("torch not installed.")
+    if not torch.cuda.is_available():
+        sys.exit("No CUDA device. Run on a GPU box.")
+    try:
+        import vllm  # noqa: F401
+    except ImportError:
+        sys.exit("vllm not installed.")
+
+
+def _get_attention_layers(vllm_config):
+    get_layers = None
+    for modpath in ("vllm.config", "vllm.model_executor.models.utils"):
+        try:
+            mod = __import__(modpath, fromlist=["get_layers_from_vllm_config"])
+            get_layers = mod.get_layers_from_vllm_config
+            break
+        except (ImportError, AttributeError):
+            continue
+    if get_layers is None:
+        raise ImportError("could not locate get_layers_from_vllm_config")
+    attention_cls = None
+    for modpath in ("vllm.model_executor.layers.attention", "vllm.attention"):
+        try:
+            attention_cls = __import__(modpath, fromlist=["Attention"]).Attention
+            break
+        except (ImportError, AttributeError):
+            continue
+    if attention_cls is None:
+        raise ImportError("could not locate the Attention layer class")
+    return get_layers(vllm_config, attention_cls)
+
+
+def _make_hook(layer_name: str):
+    def hook(module, inputs, output):
+        global _log_count
+        if _log_count >= _MAX_LOGS:
+            return
+        q_shape = None
+        if inputs and hasattr(inputs[0], "shape"):
+            q_shape = tuple(inputs[0].shape)
+        out_shape = tuple(output.shape) if hasattr(output, "shape") else None
+        print(f"[hook] {layer_name}: query={q_shape} output={out_shape}")
+        _log_count += 1
+
+    return hook
+
+
+def _install_attention_hooks() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+    if getattr(GPUModelRunner.load_model, "_h2o_probe_patched", False):
+        return
+    original_load = GPUModelRunner.load_model
+
+    def patched_load(self, *args, **kwargs):
+        result = original_load(self, *args, **kwargs)
+        try:
+            layers = _get_attention_layers(self.vllm_config)
+            for name, attn in layers.items():
+                attn.register_forward_hook(_make_hook(name))
+            print(f"[probe] installed hooks on {len(layers)} attention layers")
+        except Exception as exc:  # noqa: BLE001 - probe: surface it
+            print(f"[probe][warn] could not install hooks: {exc}")
+        return result
+
+    patched_load._h2o_probe_patched = True
+    GPUModelRunner.load_model = patched_load
+
+
+def main() -> None:
+    args = parse_args()
+    _check_env()
+    _install_attention_hooks()
+
+    from vllm import LLM, SamplingParams
+
+    llm = LLM(
+        model=args.model,
+        max_model_len=args.max_model_len,
+        enforce_eager=True,
+        enable_prefix_caching=False,
+    )
+    out = llm.generate(
+        [args.prompt], SamplingParams(max_tokens=args.max_tokens, temperature=0.0)
+    )
+    print(f"\n[probe] Output: {out[0].outputs[0].text!r}")
+    print(
+        "\nIf you saw the FLASH_ATTN... line say FLASHINFER and [hook] lines "
+        "fired, the foundation for in-engine H2O scoring is in place. Next: "
+        "capture LSE in the FlashInfer path and compute per-position score sums."
+    )
+
+
+if __name__ == "__main__":
+    main()

From edb19eb350dbb53aab77461bc654cf7a0c67c606 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 14:57:55 -0400
Subject: [PATCH 3/8] Probe: select FlashInfer via EngineArgs attention_backend
 (0.23 renamed knob)

---
 examples/vllm_h2o_probe.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/vllm_h2o_probe.py b/examples/vllm_h2o_probe.py
index 44a4888..3675d25 100644
--- a/examples/vllm_h2o_probe.py
+++ b/examples/vllm_h2o_probe.py
@@ -39,7 +39,6 @@
 import os
 
 os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
-os.environ.setdefault("VLLM_ATTENTION_BACKEND", "FLASHINFER")
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
 import sys
@@ -145,6 +144,7 @@ def main() -> None:
         max_model_len=args.max_model_len,
         enforce_eager=True,
         enable_prefix_caching=False,
+        attention_backend="FLASHINFER",  # 0.23: EngineArgs field (not an env var)
     )
     out = llm.generate(
         [args.prompt], SamplingParams(max_tokens=args.max_tokens, temperature=0.0)

From 4b1b1a6b3221406d36bd9c43127e860c5fbc1468 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 15:03:07 -0400
Subject: [PATCH 4/8] Record 4.3 step-1: FlashInfer selectable via
 attention_backend; hooks fire

---
 docs/vllm_integration.md | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/docs/vllm_integration.md b/docs/vllm_integration.md
index 6de01fa..176b36a 100644
--- a/docs/vllm_integration.md
+++ b/docs/vllm_integration.md
@@ -275,8 +275,17 @@ with `return_lse=True` (`vllm/v1/attention/backends/flashinfer.py`), yielding th
 per-query log-sum-exp. With `LSE + Q + K`, per-position weights are
 `exp(q_i·k_j·scale - lse_i)` - the same Q,K,LSE recipe keys_values' FlashInfer
 path already uses, and summable with the existing Triton score-sum. So the plan
-is: force `VLLM_ATTENTION_BACKEND=FLASHINFER`, capture LSE per layer, compute the
-per-position score-sum, and feed `H2OManager.record_block_scores`.
+is: force FlashInfer, capture LSE per layer, compute the per-position score-sum,
+and feed `H2OManager.record_block_scores`.
+
+Step 1 confirmed on box (g5.xlarge, vLLM 0.23): FlashInfer is selected via the
+`EngineArgs` field `attention_backend="FLASHINFER"` (the old
+`VLLM_ATTENTION_BACKEND` env var is gone in 0.23). Forward hooks on the 24
+`Attention` layers install and fire during generation, and output stays correct.
+Foundation harness: `examples/vllm_h2o_probe.py`. Remaining step-2 work: the
+layer forward hook sees Q/output but not LSE, which lives inside the FlashInfer
+impl's run path - so LSE capture needs a hook one level down (on the FlashInfer
+`AttentionImpl.forward` / wrapper `.run`), not the module hook.
 
 ### Blocker 2 - non-prefix eviction in the block table (the hard part)
 vLLM's built-in managers only ever free a contiguous **prefix** (sliding window,

From 4d9d41761f422eb3435923d314505f7258112f84 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 15:10:19 -0400
Subject: [PATCH 5/8] Add task 4.3 step-2a probe: inspect
 FlashInferImpl.forward args/metadata

---
 examples/vllm_h2o_lse_probe.py | 156 +++++++++++++++++++++++++++++++++
 1 file changed, 156 insertions(+)
 create mode 100644 examples/vllm_h2o_lse_probe.py

diff --git a/examples/vllm_h2o_lse_probe.py b/examples/vllm_h2o_lse_probe.py
new file mode 100644
index 0000000..4254bbf
--- /dev/null
+++ b/examples/vllm_h2o_lse_probe.py
@@ -0,0 +1,156 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Task 4.3 probe, step 2a: inspect the FlashInfer attention impl.
+
+To compute H2O scores in-engine we need the per-query LSE (or Q + paged K) and
+the attention metadata (block tables, seq lens) that maps KV positions to
+blocks. The module-level forward hook only exposed Q/output; those tensors live
+one level down, inside ``FlashInferImpl.forward``.
+
+This probe wraps ``FlashInferImpl.forward`` and logs, for the first few calls,
+the shape/dtype of every positional and keyword argument plus the relevant
+attention-metadata fields. It does not change behavior (it calls the original).
+The point is to learn exactly what we can reach before writing the LSE -> score
+computation.
+
+Usage:
+    python examples/vllm_h2o_lse_probe.py --model Qwen/Qwen2.5-0.5B-Instruct
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+import sys
+
+_MAX_LOGS = 4
+_log_count = 0
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    p.add_argument("--max-model-len", type=int, default=2048)
+    p.add_argument(
+        "--prompt",
+        default="List three uses of a key-value cache, then summarize them.",
+    )
+    p.add_argument("--max-tokens", type=int, default=8)
+    return p.parse_args()
+
+
+def _check_env() -> None:
+    try:
+        import torch
+    except ImportError:
+        sys.exit("torch not installed.")
+    if not torch.cuda.is_available():
+        sys.exit("No CUDA device. Run on a GPU box.")
+    try:
+        import vllm  # noqa: F401
+    except ImportError:
+        sys.exit("vllm not installed.")
+
+
+def _describe(x) -> str:
+    if hasattr(x, "shape") and hasattr(x, "dtype"):
+        return f"Tensor{tuple(x.shape)}/{x.dtype}"
+    if isinstance(x, (int, float, bool, str)) or x is None:
+        return repr(x)
+    return type(x).__name__
+
+
+def _describe_metadata(md) -> str:
+    """Pull out the attn-metadata fields we care about for scoring."""
+    fields = [
+        "num_prefills",
+        "num_prefill_tokens",
+        "num_decode_tokens",
+        "num_actual_tokens",
+        "slot_mapping",
+        "block_table",
+        "block_tables",
+        "seq_lens",
+        "query_start_loc",
+    ]
+    parts = []
+    for f in fields:
+        if hasattr(md, f):
+            parts.append(f"{f}={_describe(getattr(md, f))}")
+    return "; ".join(parts) if parts else f"(no known fields on {type(md).__name__})"
+
+
+def _wrap_flashinfer_impl() -> None:
+    from vllm.v1.attention.backends.flashinfer import FlashInferImpl
+
+    if getattr(FlashInferImpl.forward, "_h2o_probe_wrapped", False):
+        return
+    original = FlashInferImpl.forward
+
+    def wrapped(self, *args, **kwargs):
+        global _log_count
+        if _log_count < _MAX_LOGS:
+            print(f"\n[impl] FlashInferImpl.forward call #{_log_count}")
+            print(
+                f"  can_return_lse_for_decode = "
+                f"{getattr(self, 'can_return_lse_for_decode', None)}"
+            )
+            for i, a in enumerate(args):
+                desc = _describe(a)
+                print(f"  arg[{i}] = {desc}")
+                # The attention metadata is whichever arg looks like it.
+                if "Metadata" in type(a).__name__ or hasattr(a, "slot_mapping"):
+                    print(f"    metadata: {_describe_metadata(a)}")
+            for k, v in kwargs.items():
+                print(f"  kw[{k}] = {_describe(v)}")
+            _log_count += 1
+        return original(self, *args, **kwargs)
+
+    wrapped._h2o_probe_wrapped = True
+    FlashInferImpl.forward = wrapped
+    print("[probe] wrapped FlashInferImpl.forward")
+
+
+def main() -> None:
+    args = parse_args()
+    _check_env()
+    _wrap_flashinfer_impl()
+
+    from vllm import LLM, SamplingParams
+
+    llm = LLM(
+        model=args.model,
+        max_model_len=args.max_model_len,
+        enforce_eager=True,
+        enable_prefix_caching=False,
+        attention_backend="FLASHINFER",
+    )
+    out = llm.generate(
+        [args.prompt], SamplingParams(max_tokens=args.max_tokens, temperature=0.0)
+    )
+    print(f"\n[probe] Output: {out[0].outputs[0].text!r}")
+    print(
+        "\nWhat to look for: which arg is the query, the kv_cache, and the "
+        "attn metadata (block_table/slot_mapping/seq_lens). Those determine how "
+        "step 2b computes per-position scores and maps them to blocks."
+    )
+
+
+if __name__ == "__main__":
+    main()

From deefe5627d498d13a238ba5cef7f57b452a98c95 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 15:22:11 -0400
Subject: [PATCH 6/8] Probe: log real (non-warmup) FlashInfer forward calls +
 full metadata fields

---
 examples/vllm_h2o_lse_probe.py | 52 +++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 26 deletions(-)

diff --git a/examples/vllm_h2o_lse_probe.py b/examples/vllm_h2o_lse_probe.py
index 4254bbf..85b9617 100644
--- a/examples/vllm_h2o_lse_probe.py
+++ b/examples/vllm_h2o_lse_probe.py
@@ -77,23 +77,20 @@ def _describe(x) -> str:
 
 
 def _describe_metadata(md) -> str:
-    """Pull out the attn-metadata fields we care about for scoring."""
-    fields = [
-        "num_prefills",
-        "num_prefill_tokens",
-        "num_decode_tokens",
-        "num_actual_tokens",
-        "slot_mapping",
-        "block_table",
-        "block_tables",
-        "seq_lens",
-        "query_start_loc",
-    ]
+    """Dump the public tensor/scalar attributes of the attn metadata object."""
     parts = []
-    for f in fields:
-        if hasattr(md, f):
-            parts.append(f"{f}={_describe(getattr(md, f))}")
-    return "; ".join(parts) if parts else f"(no known fields on {type(md).__name__})"
+    for name in sorted(n for n in dir(md) if not n.startswith("_")):
+        try:
+            val = getattr(md, name)
+        except Exception:  # noqa: BLE001
+            continue
+        if callable(val):
+            continue
+        if hasattr(val, "shape") and hasattr(val, "dtype"):
+            parts.append(f"{name}={_describe(val)}")
+        elif isinstance(val, (int, float, bool)) or val is None:
+            parts.append(f"{name}={val!r}")
+    return "\n    ".join(parts) if parts else f"(nothing on {type(md).__name__})"
 
 
 def _wrap_flashinfer_impl() -> None:
@@ -105,20 +102,23 @@ def _wrap_flashinfer_impl() -> None:
 
     def wrapped(self, *args, **kwargs):
         global _log_count
-        if _log_count < _MAX_LOGS:
-            print(f"\n[impl] FlashInferImpl.forward call #{_log_count}")
+        # args: (layer, query, key, value, kv_cache, attn_metadata)
+        attn_metadata = args[5] if len(args) > 5 else None
+        # Skip warmup/profiling calls (no metadata, empty kv_cache).
+        if attn_metadata is not None and _log_count < _MAX_LOGS:
+            kv_cache = args[4] if len(args) > 4 else None
+            query, key, value = args[1], args[2], args[3]
+            print(f"\n[impl] real FlashInferImpl.forward call #{_log_count}")
             print(
                 f"  can_return_lse_for_decode = "
                 f"{getattr(self, 'can_return_lse_for_decode', None)}"
             )
-            for i, a in enumerate(args):
-                desc = _describe(a)
-                print(f"  arg[{i}] = {desc}")
-                # The attention metadata is whichever arg looks like it.
-                if "Metadata" in type(a).__name__ or hasattr(a, "slot_mapping"):
-                    print(f"    metadata: {_describe_metadata(a)}")
-            for k, v in kwargs.items():
-                print(f"  kw[{k}] = {_describe(v)}")
+            print(f"  query = {_describe(query)}")
+            print(f"  key   = {_describe(key)}")
+            print(f"  value = {_describe(value)}")
+            print(f"  kv_cache = {_describe(kv_cache)}")
+            print(f"  attn_metadata ({type(attn_metadata).__name__}):")
+            print(f"    {_describe_metadata(attn_metadata)}")
             _log_count += 1
         return original(self, *args, **kwargs)
 

From 65716ea5b692cebcc56ab11d8c8cb287a607c836 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 15:32:13 -0400
Subject: [PATCH 7/8] Record 4.3 step-2a findings: FlashInfer impl surface,
 kv_cache layout, effort note

---
 docs/vllm_integration.md | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/docs/vllm_integration.md b/docs/vllm_integration.md
index 176b36a..b574f93 100644
--- a/docs/vllm_integration.md
+++ b/docs/vllm_integration.md
@@ -301,13 +301,29 @@ append-only. Candidate approaches, to evaluate empirically:
 - **B. Dense-cache attention backend (Option B)**: faithful per-head eviction,
   most invasive.
 
-### Recommended 4.3 order
-1. Wire blocker 1 first (LSE capture -> score-sum -> `record_block_scores`) and
-   verify scores look sane during generation; this is reusable regardless of how
-   eviction is wired.
+### FlashInfer impl surface (mapped on box, vLLM 0.23)
+`FlashInferImpl.forward(self, layer, query, key, value, kv_cache, attn_metadata,
+*, output, ...)`:
+- `query (T, n_heads, head_size)`, `key/value (T, n_kv_heads, head_size)` for the
+  current step's `T` tokens (a batch spanning prefill + decode requests).
+- `kv_cache (num_blocks, 2, block_size, n_kv_heads, head_size)` - paged K/V.
+- `attn_metadata` (`FlashInferMetadata`): `num_decodes`, `num_decode_tokens`,
+  `num_prefills`, `num_prefill_tokens`, `slot_mapping (T,)`, `use_cascade`.
+  The block table is not a top-level field; it lives in the nested decode/prefill
+  FlashInfer wrapper objects. `can_return_lse_for_decode = True`.
+
+### Recommended 4.3 order and effort
+1. Wire blocker 1 (LSE -> score-sum -> `record_block_scores`). The per-request
+   block mapping comes from `H2OManager.req_to_blocks` (we own it); the hard part
+   is correlating each of the `T` batched tokens/scores back to its request and
+   layer, every step.
 2. Prototype approach A (compaction) on short sequences; compare against the
-   keys_values/LitGPT H2O reference. If compaction can't be made consistent with
-   V1 metadata, escalate to Option B.
+   keys_values/LitGPT H2O reference; escalate to Option B if inconsistent.
+
+A correct in-engine H2O scoring+eviction loop is a multi-day effort (per-request
+score routing across the batched FlashInfer call, plus the unsolved
+arbitrary-eviction "holes" problem). De-risked so far: weights reachable
+(LSE / Q+paged-K), FlashInfer backend works, hooks fire, impl surface mapped.
 
 ## Open questions
 

From 8f62a7a3a6779b39a4aa9ec403937103676699c2 Mon Sep 17 00:00:00 2001
From: Michael Jerge <112141470+mmjerge@users.noreply.github.com>
Date: Wed, 24 Jun 2026 16:17:02 -0400
Subject: [PATCH 8/8] Add task 4.3 step-2b.1: in-engine per-position H2O scores
 (single-seq decode)

---
 examples/vllm_h2o_score_probe.py | 177 +++++++++++++++++++++++++++++++
 1 file changed, 177 insertions(+)
 create mode 100644 examples/vllm_h2o_score_probe.py

diff --git a/examples/vllm_h2o_score_probe.py b/examples/vllm_h2o_score_probe.py
new file mode 100644
index 0000000..c792f42
--- /dev/null
+++ b/examples/vllm_h2o_score_probe.py
@@ -0,0 +1,177 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Task 4.3 step 2b.1: compute per-KV-position H2O scores in-engine (single seq).
+
+Smallest real scoring step. For a single-sequence generation, we wrap
+``FlashInferImpl.forward`` and, per layer, accumulate the ``slot_mapping`` across
+calls. For one request that enumerates exactly its KV slots in position order, so
+we can gather the request's full K/V from the paged ``kv_cache`` and compute
+per-position attention mass with the verified ``reference_summed_attention``
+(``keys_values.vllm.attention``) - no block table or manager wiring needed yet.
+
+For the first few decode steps of layer 0, it prints the score vector length,
+its total mass (a sanity check: ~= number of query heads per KV group), and the
+top-scoring positions. This validates the in-engine score signal before we
+tackle multi-request batching, the block mapping, and eviction.
+
+Run single-sequence only (one prompt). Usage:
+    python examples/vllm_h2o_score_probe.py --model Qwen/Qwen2.5-0.5B-Instruct
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+import sys
+from collections import defaultdict
+
+# Accumulated KV slots per layer-impl instance (id(self) -> list[int]).
+_slots_by_layer: dict = defaultdict(list)
+_decode_logs = 0
+_MAX_DECODE_LOGS = 4
+_first_layer_id = None
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    p.add_argument("--max-model-len", type=int, default=2048)
+    p.add_argument(
+        "--prompt",
+        default="The capital of France is Paris. Tell me a short fact about it.",
+    )
+    p.add_argument("--max-tokens", type=int, default=12)
+    return p.parse_args()
+
+
+def _check_env() -> None:
+    try:
+        import torch  # noqa: F401
+    except ImportError:
+        sys.exit("torch not installed.")
+    import torch
+
+    if not torch.cuda.is_available():
+        sys.exit("No CUDA device. Run on a GPU box.")
+    try:
+        import vllm  # noqa: F401
+    except ImportError:
+        sys.exit("vllm not installed.")
+
+
+def _score_decode_step(self, query, kv_cache, slots) -> None:
+    """Gather the request's K/V from the paged cache and print H2O scores."""
+    import torch
+
+    from keys_values.vllm.attention import reference_summed_attention
+
+    # kv_cache: (num_blocks, 2, block_size, n_kv_heads, head_size)
+    num_blocks, two, block_size, n_kv_heads, head_size = kv_cache.shape
+    keys_all = kv_cache[:, 0].reshape(-1, n_kv_heads, head_size)
+    values_all = kv_cache[:, 1].reshape(-1, n_kv_heads, head_size)
+    slot_idx = torch.tensor(slots, device=kv_cache.device, dtype=torch.long)
+    k_req = keys_all[slot_idx]  # (seq_len, n_kv_heads, head_size)
+    v_req = values_all[slot_idx]
+    seq_len = k_req.shape[0]
+
+    # Shape for the reference: query (1, n_heads, 1, head_size);
+    # key/value (1, n_kv_heads, seq_len, head_size).
+    n_heads = query.shape[1]
+    q = query.reshape(1, n_heads, 1, head_size)
+    k = k_req.permute(1, 0, 2).unsqueeze(0)
+    v = v_req.permute(1, 0, 2).unsqueeze(0)
+    _, summed = reference_summed_attention(q, k, v, causal=True)
+    summed = summed[0]  # (n_kv_heads, seq_len)
+    per_head_mass = summed.sum(dim=-1)  # ~= query heads per group
+    top = torch.topk(summed.sum(dim=0), k=min(3, seq_len))
+    print(
+        f"  [score] seq_len={seq_len} mass_per_kv_head={per_head_mass.tolist()} "
+        f"top_positions={top.indices.tolist()} top_scores="
+        f"{[round(x, 3) for x in top.values.tolist()]}"
+    )
+
+
+def _wrap_flashinfer_impl() -> None:
+    from vllm.v1.attention.backends.flashinfer import FlashInferImpl
+
+    if getattr(FlashInferImpl.forward, "_h2o_score_wrapped", False):
+        return
+    original = FlashInferImpl.forward
+
+    def wrapped(self, *args, **kwargs):
+        global _decode_logs, _first_layer_id
+        attn_metadata = args[5] if len(args) > 5 else None
+        if attn_metadata is not None:
+            query = args[1]
+            kv_cache = args[4]
+            slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+            if (
+                kv_cache is not None
+                and kv_cache.numel() > 0
+                and slot_mapping is not None
+            ):
+                lid = id(self)
+                if _first_layer_id is None:
+                    _first_layer_id = lid
+                _slots_by_layer[lid].extend(slot_mapping.tolist())
+                num_decode = int(getattr(attn_metadata, "num_decode_tokens", 0))
+                num_prefills = int(getattr(attn_metadata, "num_prefills", 0))
+                is_single_decode = num_prefills == 0 and num_decode == 1
+                if (
+                    is_single_decode
+                    and lid == _first_layer_id
+                    and _decode_logs < _MAX_DECODE_LOGS
+                ):
+                    print(f"\n[impl] layer0 decode step (log {_decode_logs}):")
+                    _score_decode_step(self, query, kv_cache, _slots_by_layer[lid])
+                    _decode_logs += 1
+        return original(self, *args, **kwargs)
+
+    wrapped._h2o_score_wrapped = True
+    FlashInferImpl.forward = wrapped
+    print("[probe] wrapped FlashInferImpl.forward for in-engine scoring")
+
+
+def main() -> None:
+    args = parse_args()
+    _check_env()
+    _wrap_flashinfer_impl()
+
+    from vllm import LLM, SamplingParams
+
+    llm = LLM(
+        model=args.model,
+        max_model_len=args.max_model_len,
+        enforce_eager=True,
+        enable_prefix_caching=False,
+        attention_backend="FLASHINFER",
+    )
+    out = llm.generate(
+        [args.prompt], SamplingParams(max_tokens=args.max_tokens, temperature=0.0)
+    )
+    print(f"\n[probe] Output: {out[0].outputs[0].text!r}")
+    print(
+        "\nSanity check: mass_per_kv_head should be ~= query heads per KV group "
+        "(14/2 = 7 for Qwen2.5-0.5B), and seq_len should grow by 1 each decode "
+        "step. If so, we have a correct in-engine H2O score signal."
+    )
+
+
+if __name__ == "__main__":
+    main()