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

docs/vllm_integration.md

@@ -349,9 +349,72 @@ wrappers with `return_lse=True`, so the log-sum-exp needed to normalize is
 available on the same call. So we neither need a vLLM feature request nor a
 second attention pass; the remaining work is purely *wiring* this into the
 in-engine attention path and mapping per-position scores to blocks (prototyped
-in `examples/vllm_h2o_score_probe.py` on the follow-up branch — gathers a
-request's paged K/V and computes correct per-position mass in-engine for the
-single-sequence decode case).
+in `examples/vllm_h2o_score_probe.py` — gathers a request's paged K/V and
+computes correct per-position mass in-engine for the single-sequence decode
+case).
+
+## H2O wiring (task 4.3): findings and plan
+
+Two sub-problems must be solved to run H2O in-engine; investigation status.
+(Blocker 1 below is the *in-engine wiring* of the weight signal whose math is
+already implemented in the task-3 module above.)
+
+### 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 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,
+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.
+
+### 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; 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
 

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:
+    """Dump the public tensor/scalar attributes of the attn metadata object."""
+    parts = []
+    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:
+    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
+        # 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)}"
+            )
+            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)
+
+    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()

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_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,
+        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)
+    )
+    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()