diff --git a/docs/vllm_integration.md b/docs/vllm_integration.md index a8c1bfd..f2a052d 100644 --- a/docs/vllm_integration.md +++ b/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 diff --git a/examples/vllm_h2o_lse_probe.py b/examples/vllm_h2o_lse_probe.py new file mode 100644 index 0000000..85b9617 --- /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: + """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() diff --git a/examples/vllm_h2o_probe.py b/examples/vllm_h2o_probe.py new file mode 100644 index 0000000..3675d25 --- /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_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() 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()