flow: add H2O sparse attention (page-level and token-level, MHA + GQA)

[GeoffreyWang1117]

Four new flows implementing the H2O heavy-hitter criterion (arXiv:2306.14048), registered as h2o_sparse_attention, gqa_h2o_sparse_attention, h2o_token_sparse_attention, gqa_h2o_token_sparse_attention. No existing flow is touched; the only change outside flow/algorithms.py is a one-line export.

Page-level variants

forward_cache runs once per page and cannot accumulate across decode steps, so these rank pages by two per-page statistics:

• average relevance — q · mean(k), i.e. the page's total pre-softmax attention mass up to the constant page size;

• peak relevance — an upper bound on max_k q·k from the element-wise max/min envelope:

  max_k q·k  ≤  Σ_d max(q_d · max_d,  q_d · min_d)
  

Both envelopes are required. With a signed query, q_d · max_d alone is an upper bound only where q_d ≥ 0, and is a lower bound everywhere else. This is the same bound gqa_quest_sparse_attention already uses, applied to the heavy-hitter score.

The two terms are not on the same scale — the peak term is an upper bound over the page and is systematically larger than the mean term — so they are combined through Add(alpha=w_avg, beta=w_peak) with the weights exposed on the constructor rather than fixed at 1.

Token-level variants

With page_size=1 a page is a token, and a Save inside forward_indexer carries the running score across decode steps (forward_cache cannot). Three points worth review:

1. The accumulated quantity is the attention probability, not the raw logit. A Softmax(dim=0) over the request's tokens is applied before accumulation. H2O accumulates normalised attention mass; a running sum of raw logits is unbounded and sign-indefinite. The softmax scale defaults to head_dim ** -0.5, resolved in create_cache (the only hook that receives head_dim) and overridable via the constructor.

2. forward_cache zeroes cache["hh_score"] for the pages written in the step. The KV pool is torch.zeros once at allocation and never again; KVCacheAllocator.free() only returns indices to the free list. A page released by a finished request keeps its contents and is handed to the next request as-is, so without this reset a new request inherits the previous one's accumulated heavy-hitter scores. The page-level flows are immune because CMean/CMax/CMin are overwriting reductions.

This is what the new Fill export is for — vortex_torch/cache/fill.py already implements it (_impl_map = {FORMAT.PAGED: fill_p}), it just was not re-exported from the package. fill_p_kernel triggers on end-of-page tokens, which with page_size=1 is every token, so it resets exactly the pages written this step.

3. The accumulator decays (default 0.984375 = 1 − 2⁻⁶). The KV pool is bfloat16 (assert self.store_dtype == torch.bfloat16), so hh_score accumulates in bfloat16 (~8 bits of mantissa). An undecayed sum grows without bound; once it does, each increment is small relative to the accumulated value and is lost to rounding, freezing the ranking at whatever it was. A decay keeps the accumulator at a bounded steady state of 1/(1−decay) typical increments, which keeps them representable. 0.984375 is exactly representable in bfloat16 and, in simulation, the longest window that still reproduces the top-k of an exact float64 accumulator. decay=1.0 restores the plain sum for anyone who wants it.

page_size == 1 is asserted in create_cache. disable_radix_cache=True is also required — prefix sharing would let one request's scores leak into another through shared pages — but that cannot be checked from inside a flow, so it is documented on the class instead.

Verification

RTX 3090, torch 2.9.1, triton 3.5.1.

Profile mode — real op dispatch, shape/format validation and buffer allocation, for all four flows across group sizes 1, 3, 4, 7, 8, plus the page_size guard and the resolved softmax scale:

  [ok ] h2o_sparse_attention             G=1 page_size=16
  [ok ] h2o_sparse_attention             G=3 page_size=16
  [ok ] h2o_sparse_attention             G=4 page_size=16
  [ok ] h2o_token_sparse_attention       G=1 page_size=1
  [ok ] h2o_token_sparse_attention       G=3 page_size=1
  [ok ] gqa_h2o_token_sparse_attention   G=1 page_size=1
  [ok ] gqa_h2o_token_sparse_attention   G=3 page_size=1

page_size guard on the token-level flows:
  [ok ] h2o_token_sparse_attention: rejected page_size=16
  [ok ] gqa_h2o_token_sparse_attention: rejected page_size=16

resolved softmax scale (expect head_dim**-0.5 = 0.088388):
  [ok ] h2o_token_sparse_attention: scale=0.088388
  [ok ] h2o_token_sparse_attention: explicit scale honoured (0.09)
  [ok ] gqa_h2o_token_sparse_attention: scale=0.088388
  [ok ] gqa_h2o_token_sparse_attention: explicit scale honoured (0.09)

Add(alpha, beta) decay wiring:
  [ok ] h2o_token_sparse_attention: alpha=1.0 beta=0.984375
  [ok ] h2o_token_sparse_attention: decay=1.0 -> beta=1.0
  [ok ] gqa_h2o_token_sparse_attention: alpha=1.0 beta=0.984375
  [ok ] gqa_h2o_token_sparse_attention: decay=1.0 -> beta=1.0

all checks passed

Execute mode — the actual Triton kernels behind Fill, Add and Save (topK needs the compiled CUDA extension and is not what these fixes touch):

Fill(0.0) resets hh_score for exactly the pages written this step
  [ok ] touched pages zeroed  values=[0.0, 0.0, 0.0, 0.0]
  [ok ] untouched pages preserved

Add(alpha=1, beta=decay) + Save recurrence, 128 tokens x 800 steps
  decay=0.984375  (default               ) rel.err=8.40e-03  top-32 vs exact=100%
  [ok ] default decay tracks the exact accumulator
  [ok ] accumulator stays bounded  max=18.2
  decay=1.0       (decay=1.0 (undecayed) ) rel.err=4.17e-01  top-32 vs exact=97%
  decay=0.984375  response to a changed step 800: max delta = 1.562e-01
  [ok ] still responds after 800 steps
  decay=1.0       response to a changed step 800: max delta = 2.500e-01

all checks passed

The undecayed row is the comparison that motivates the default: 42% relative error against the exact accumulator after 800 steps, versus 0.84% with the decay.

Verification scripts

Profile mode:

"""Profile-mode smoke test for the H2O flows.

Drives forward_indexer through vFlow.run_indexer_virtual (real op dispatch,
shape/format validation, buffer allocation) and forward_cache through the same
path the memory pool uses at startup. Also checks the page_size guard and the
resolved softmax scale.
"""
import torch
import vortex_torch.flow.algorithms as A
from vortex_torch.cache import Context as CContext
from vortex_torch.abs import as_vtensor, FORMAT

HEAD_DIM = 128
FLOWS = [
    ("h2o_sparse_attention", A.H2OSparseAttention, 16),
    ("gqa_h2o_sparse_attention", A.GQAH2OSparseAttention, 16),
    ("h2o_token_sparse_attention", A.H2OTokenSparseAttention, 1),
    ("gqa_h2o_token_sparse_attention", A.GQAH2OTokenSparseAttention, 1),
]


def profile_cache(flow, page_size, group_size):
    """Mirror vtx_graph_memory_pool._initialize_graph's forward_cache profiling."""
    ctx = CContext()
    ctx.head_num = 1
    ctx.page_size = page_size
    meta = flow.get_cache_meta_info(page_size, HEAD_DIM)
    cache = {
        name: as_vtensor(
            torch.zeros((0, s[0], s[1]), dtype=torch.bfloat16, device="cuda"),
            FORMAT.PAGED,
        )
        for name, s in meta.items()
    }
    loc = torch.empty((0,), dtype=torch.int64, device="cuda")
    flow.forward_cache(cache=cache, loc=loc, ctx=ctx)


fails = 0
for name, cls, page_size in FLOWS:
    for group_size in (1, 3, 4, 7, 8):
        try:
            f = cls()
            f.get_cache_meta_info(page_size, HEAD_DIM)   # resolves softmax scale
            f.run_indexer_virtual(group_size, page_size, HEAD_DIM)
            profile_cache(cls(), page_size, group_size)
            print(f"  [ok ] {name:<32} G={group_size} page_size={page_size}")
        except Exception as e:
            fails += 1
            m = str(e).strip().splitlines()
            print(f"  [FAIL] {name:<32} G={group_size}: {type(e).__name__}: "
                  f"{m[-1][:120] if m else e}")

print("\npage_size guard on the token-level flows:")
for name, cls, _ in FLOWS[2:]:
    try:
        cls().create_cache(page_size=16, head_dim=HEAD_DIM)
        print(f"  [FAIL] {name}: page_size=16 was accepted")
        fails += 1
    except AssertionError as e:
        print(f"  [ok ] {name}: rejected page_size=16")

print("\nresolved softmax scale (expect head_dim**-0.5 = "
      f"{HEAD_DIM ** -0.5:.6f}):")
for name, cls, ps in FLOWS[2:]:
    f = cls(); f.create_cache(ps, HEAD_DIM)
    ok = abs(f.softmax.scale - HEAD_DIM ** -0.5) < 1e-9
    fails += 0 if ok else 1
    print(f"  [{'ok ' if ok else 'FAIL'}] {name}: scale={f.softmax.scale:.6f}")
    f2 = cls(scale=0.09); f2.create_cache(ps, HEAD_DIM)
    ok2 = abs(f2.softmax.scale - 0.09) < 1e-9
    fails += 0 if ok2 else 1
    print(f"  [{'ok ' if ok2 else 'FAIL'}] {name}: explicit scale honoured "
          f"({f2.softmax.scale})")

print("\nAdd(alpha, beta) decay wiring:")
for name, cls, ps in FLOWS[2:]:
    f = cls()
    ok = f.add.alpha == 1.0 and abs(f.add.beta - 0.984375) < 1e-12
    fails += 0 if ok else 1
    print(f"  [{'ok ' if ok else 'FAIL'}] {name}: alpha={f.add.alpha} beta={f.add.beta}")
    f = cls(decay=1.0)
    ok = f.add.beta == 1.0
    fails += 0 if ok else 1
    print(f"  [{'ok ' if ok else 'FAIL'}] {name}: decay=1.0 -> beta={f.add.beta}")

print(f"\n{'FAILURES: ' + str(fails) if fails else 'all checks passed'}")
raise SystemExit(1 if fails else 0)

Execute mode:

"""Execute-mode verification of the two H2O token-level fixes.

Runs the real Triton kernels behind the ops the flow uses:
  Fill  -> _fill_p                     (fix 1: reset hh_score on page reuse)
  Add   -> _elementwise_binary_rpr     (fix 2: attn + decay * old_score)
  Save  -> _save_rp                    (persist back to the paged cache)

topK is skipped (needs the compiled CUDA extension); it is not what the two
fixes touch.
"""
import torch
from vortex_torch.cache.triton_kernels.fill_impl import _fill_p
from vortex_torch.indexer.triton_kernels.elementwise_binary_impl import _elementwise_binary_rpr
from vortex_torch.indexer.triton_kernels.save_load_impl import _save_rp
from vortex_torch.utils import ElementwiseBinaryOpType

DEV, CHUNK, NSMS = "cuda", 16, 8
fails = []


def winfo(n, chunk=CHUNK):
    offs, lens, o = [], [], 0
    while o < n:
        L = min(chunk, n - o)
        offs.append(o); lens.append(L); o += L
    t = lambda v: torch.tensor(v, dtype=torch.int32, device=DEV)
    return t(offs), t(lens), t([len(offs)])


def check(name, ok, detail=""):
    print(f"  [{'ok ' if ok else 'FAIL'}] {name}{('  ' + detail) if detail else ''}")
    if not ok:
        fails.append(name)


# ---------------------------------------------------------------- fix 1
print("Fill(0.0) resets hh_score for exactly the pages written this step")
NUM_PAGES = 64
pool = torch.full((NUM_PAGES, 1, 1), 7.5, device=DEV, dtype=torch.bfloat16)  # stale
touched = torch.tensor([3, 5, 11, 40], dtype=torch.int64, device=DEV)
_fill_p(pool, touched, num_kv_heads=1, page_size=1, alpha=0.0)
torch.cuda.synchronize()
flat = pool.view(-1)
mask = torch.zeros(NUM_PAGES, dtype=torch.bool, device=DEV)
mask[touched] = True
check("touched pages zeroed", bool((flat[mask] == 0).all()),
      f"values={flat[mask].tolist()}")
check("untouched pages preserved", bool((flat[~mask] == 7.5).all()))


# ---------------------------------------------------------------- fix 2
def accumulate(steps, decay, probs, num_pages):
    """Drive Add + Save exactly as forward_indexer does, one step at a time."""
    hh = torch.zeros((num_pages, 1, 1), device=DEV, dtype=torch.bfloat16)
    idx = torch.arange(num_pages, dtype=torch.int32, device=DEV)   # ragged pos -> page
    offs, lens, nw = winfo(num_pages)
    out = torch.empty_like(hh)
    for s in range(steps):
        attn = probs[s].view(num_pages, 1, 1).to(torch.bfloat16).contiguous()
        _elementwise_binary_rpr(attn, hh, out, idx, offs, lens, nw, CHUNK,
                                ElementwiseBinaryOpType.Add, 1.0, decay, NSMS)
        _save_rp(out, hh, idx, offs, lens, nw, CHUNK, NSMS)
    torch.cuda.synchronize()
    return hh.view(-1).float()


def reference(steps, decay, probs):
    """Same recurrence in float64, on the bf16-rounded decay constant."""
    d = float(torch.tensor(decay, dtype=torch.bfloat16))
    acc = torch.zeros(probs.shape[1], dtype=torch.float64)
    for s in range(steps):
        acc = probs[s].double() + d * acc
    return acc


N_PAGES, N_STEP = 128, 800
torch.manual_seed(0)
affinity = torch.randn(N_PAGES) * 1.5
g = torch.Generator().manual_seed(1)
probs = torch.stack([
    torch.softmax(affinity + torch.randn(N_PAGES, generator=g) * 0.5, 0)
    for _ in range(N_STEP)
]).to(DEV)

print(f"\nAdd(alpha=1, beta=decay) + Save recurrence, "
      f"{N_PAGES} tokens x {N_STEP} steps")
for decay, label in ((0.984375, "default"), (1.0, "decay=1.0 (undecayed)")):
    got = accumulate(N_STEP, decay, probs, N_PAGES)
    ref = reference(N_STEP, decay, probs.cpu())
    rel = float((got.cpu().double() - ref).abs().max() / ref.abs().max())
    k = 32
    overlap = len(set(got.topk(k).indices.tolist())
                  & set(ref.topk(k).indices.tolist())) / k
    print(f"  decay={decay:<9} ({label:<22}) rel.err={rel:8.2e}  top-{k} vs exact={overlap:.0%}")
    if decay == 0.984375:
        check("default decay tracks the exact accumulator", rel < 0.10 and overlap == 1.0)
        check("accumulator stays bounded", float(got.max()) < 200,
              f"max={float(got.max()):.1f}")

# does the accumulator still respond at the end of a long run?
for decay in (0.984375, 1.0):
    hh_a = accumulate(N_STEP, decay, probs, N_PAGES)
    probs2 = probs.clone()
    probs2[-1] = probs2[-1] * 0 + 1.0 / N_PAGES     # a different final step
    hh_b = accumulate(N_STEP, decay, probs2, N_PAGES)
    moved = float((hh_a - hh_b).abs().max())
    print(f"  decay={decay:<9} response to a changed step {N_STEP}: "
          f"max delta = {moved:.3e}")
    if decay == 0.984375:
        check(f"still responds after {N_STEP} steps", moved > 1e-4)

print(f"\n{'FAILED: ' + ', '.join(fails) if fails else 'all checks passed'}")
raise SystemExit(1 if fails else 0)

Notes

The token-level design — using Save inside forward_indexer for cross-step accumulation, together with page_size=1 and a disabled radix cache — was suggested by Zhuoming Chen, and is credited in the class docstrings.

Independent of this PR: the GQA variants use GeMM/Softmax over the query-group axis, which today requires the group size to be a power of two (#5 lifts that). On v1 as-is these flows work for power-of-two group sizes, same as every other flow in the file.

[GeoffreyWang1117]

Follow-up: I got the environment for the full stack working again, so the token-level flow is now verified on hardware against a reference, rather than in profile mode plus a kernel-level simulation as in the description above.

The harness drives h2o_token_sparse_attention through the real ops — the compiled topk_output and the real sglang_plan_decode workload planner — for 600 decode steps, reads cache["hh_score"] straight out of the paged pool, and compares against a float64 recurrence computed independently. No sglang server, no model weights. RTX 3090, torch 2.7.1+cu126, CUDA 12.6.

Result

device: NVIDIA GeForce RTX 3090
flow: decay=0.984375  softmax scale=0.088388 (head_dim**-0.5 = 0.088388)

total probability mass in hh_score: 64.50 (steady state 1/(1-decay) = 64.00)
  mass predicted by the recurrence = 64.00 -> -0.8% lost to bf16 rounding

hh_score vs an independent float64 reference of the same recurrence
  step    rel.err   max |hh|   top-k vs ref  boundary/ulp
     1   2.37e-03       0.02          100%           0.8
   151   1.97e-02       0.88          100%           0.8
   301   1.83e-02       1.00           97%           0.0
   600   1.82e-02       0.97          100%           1.3

page reuse: forward_cache on recycled pages
  hh_score before reset: max=0.973 nonzero=116/128
  hh_score after  reset: max=0.000 nonzero=0/128

  [PASS] accumulates softmax probabilities (matches fp64 reference)
  [PASS] accumulator retains its probability mass under bf16
  [PASS] forward_cache resets hh_score on page reuse

ALL CHECKS PASSED

The line worth reading is the mass. Every step injects exactly 1.0 of probability mass, so a decayed sum has a closed-form steady state of 1/(1-decay). Measured 64.50 against a predicted 64.00 — that is direct evidence the accumulated quantity really is a normalised attention distribution and not a raw logit sum, which is the change described in point 1 above.

The undecayed counterfactual

Same harness, --decay 1.0, i.e. the plain sum:

device: NVIDIA GeForce RTX 3090
flow: decay=1.0  softmax scale=0.088388 (head_dim**-0.5 = 0.088388)

total probability mass in hh_score: 445.83 (undecayed: grows by ~1.0 per step, without bound)
  mass predicted by the recurrence = 600.00 -> +25.7% lost to bf16 rounding

hh_score vs an independent float64 reference of the same recurrence
  step    rel.err   max |hh|   top-k vs ref  boundary/ulp
     1   2.37e-03       0.02          100%           0.8
   151   4.03e-02       2.33           97%           1.1
   301   1.07e-01       4.78           97%           0.2
   600   3.19e-01       8.00           62%           0.5

page reuse: forward_cache on recycled pages
  hh_score before reset: max=8.000 nonzero=116/128
  hh_score after  reset: max=0.000 nonzero=0/128

  [FAIL] accumulates softmax probabilities (matches fp64 reference)
  [FAIL] accumulator retains its probability mass under bf16
  [PASS] forward_cache resets hh_score on page reuse

FAILED: 26% of the probability mass lost to rounding (retained 445.8 of 600.0); final rel.err 3.19e-01 >= 0.15

After 600 steps it has retained 445.8 of the 600.0 of mass the recurrence predicts — 26% rounded away — the error against the reference has grown to 3.2e-01, and top-k agreement has fallen to 62%. max |hh| = 8.00 is exact, which is the tell: bf16 near 8.0 has an ulp of 0.03125, and the per-step increments (~1/116 on average) are far below half of that, so they round to nothing. The accumulator has frozen. That is the failure mode the decay was introduced to avoid, and it reproduces on hardware.

Two things I had to fix in the harness first, in case they are useful

The scored window is not the whole sequence. My first reference softmaxed over all seq_len tokens and reported a 92% relative error, which looked like a bug in the flow. It was the reference: hh_score is nonzero exactly on indices [page_reserved_bos, seq_len - page_reserved_eos) — 116 of 128 with bos=4, eos=8 — because softmax and top-k only operate between the reserved pages, which are kept unconditionally. Normalising over that window instead drops the error to 3.2e-03. Worth knowing for anyone else writing a check against these flows.

Top-k agreement is meaningless without heavy hitters. With i.i.d. keys and queries every token is attended about equally: the 116 scores landed in 0.49–0.62 and the gap at the top-32 boundary was 0.000079, which is 28× smaller than one bf16 ulp at that magnitude. No implementation can resolve that ordering. The harness now uses keys with a stable per-token affinity along a shared direction the queries lean toward — which is the situation H2O exists for — and prints the boundary width in ulps next to each overlap, so a low number is interpretable rather than alarming. The 97% row above sits at a 0.0-ulp boundary and is correctly not treated as a failure.

Script

verify_h2o_oracle.py

"""Oracle verification for the token-level H2O flow.

Drives `h2o_token_sparse_attention` through the real indexer ops -- including
the compiled `topk_output` and the real `sglang_plan_decode` workload planner --
over a sequence of decode steps, and checks the accumulated `hh_score` against a
torch reference computed independently.

This is the check the flow's three correctness properties actually need, none of
which is observable from generated text:

  1. `hh_score` accumulates *attention probabilities* (softmax over the request's
     tokens), not raw logits.
  2. The accumulator decays, so it stays bounded and keeps responding to new
     attention after hundreds of steps instead of freezing under bf16 rounding.
  3. `forward_cache` zeroes `hh_score` for pages written this step, so a page
     recycled from a finished request does not carry its scores into the next.

Requires: built `vortex_torch_C`, one CUDA device. No sglang, no model weights.

    python experiments/verify_h2o_oracle.py
"""

import argparse
import math

import torch

from vortex_torch.abs import FORMAT, as_vtensor
from vortex_torch.flow.algorithms import H2OTokenSparseAttention
from vortex_torch.cache import Context as CacheContext
from vortex_torch.indexer import Context as IContext
from vortex_torch.indexer.utils_sglang import plan_decode
from vortex_torch.utils import Mode

DEV = "cuda"
DT = torch.bfloat16


# --------------------------------------------------------------------------
# context construction
# --------------------------------------------------------------------------
def build_ctx(max_reqs, max_pages_per_req, num_kv_heads, group_size, head_dim,
              topk_val, bos, eos, max_chunk=64, min_chunk=1):
    ctx = IContext()
    ctx.name = "h2o-oracle"
    ctx.mode = Mode.profile
    ctx._created = True

    ctx.group_size = group_size
    ctx.num_kv_heads = num_kv_heads
    ctx.num_qo_heads = num_kv_heads * group_size
    ctx.head_dim = head_dim
    ctx.page_size = 1
    ctx.indexer_dtype = DT

    ctx.max_num_pages_per_request = max_pages_per_req
    ctx.max_num_pages = max_pages_per_req * max_reqs * num_kv_heads
    ctx.topk_val = topk_val
    ctx.page_reserved_bos = bos
    ctx.page_reserved_eos = eos
    ctx.max_chunk_size = max_chunk
    ctx.min_chunk_size = min_chunk
    ctx.num_sms = torch.cuda.get_device_properties(0).multi_processor_count

    n_seg = max_reqs * num_kv_heads
    z = lambda n, dt=torch.int32: torch.zeros((n,), dtype=dt, device=DEV)
    ctx.dense_kv_indices = z(ctx.max_num_pages)
    ctx.sparse_kv_indices = z(ctx.max_num_pages)
    ctx.dense_kv_indptr = z(n_seg + 1)
    ctx.sparse_kv_indptr = z(n_seg + 1)
    ctx.kv_last_page_len = z(n_seg)
    ctx.batch_size = max_reqs

    ctx.max_num_workloads = (ctx.max_num_pages // max(1, min_chunk)) + n_seg
    ctx.winfo_q_indices = z(ctx.max_num_workloads)
    ctx.winfo_kv_offsets = z(ctx.max_num_workloads)
    ctx.winfo_kv_lens = z(ctx.max_num_workloads)
    ctx.winfo_num_workloads = z(1)
    ctx.winfo_chunk_size = z(1)
    ctx._aux_total_bytes = 0
    ctx._aux_total_flops = 0
    return ctx


def build_cache_ctx(num_kv_heads, head_dim, total_pages):
    """forward_cache runs against the cache-side Context, not the indexer one."""
    cctx = CacheContext()
    cctx.name = "h2o-oracle-cache"
    cctx.mode = Mode.profile
    cctx._created = True
    cctx.page_size = 1
    cctx.head_num = num_kv_heads
    cctx.head_dim = head_dim
    cctx.total_num_pages = total_pages
    cctx.max_new_tokens_per_batch = 4096
    cctx._aux_total_bytes = 0
    cctx._aux_total_flops = 0
    return cctx


def page_ids(token_slots, num_kv_heads, head):
    """VTX page index for (token slot, kv head) at page_size=1."""
    return token_slots * num_kv_heads + head


# --------------------------------------------------------------------------
# reference
# --------------------------------------------------------------------------
def reference_step(prev, q_mean, k, scale, decay, bos, eos):
    """One step of the flow's own recurrence, in float64.

    The indexer never scores the reserved pages: `page_reserved_bos` leading and
    `page_reserved_eos` trailing pages are kept unconditionally, and softmax /
    top-k operate only on the window between them. So the reference normalises
    over [bos, S-eos) and leaves the reserved entries at zero.

        score[bos:S-eos] = softmax(scale * q.k over that window) + decay * prev
    """
    n = k.shape[0]
    logits = (k.double() @ q_mean.double()) * scale
    probs = torch.zeros(n, dtype=torch.float64, device=k.device)
    probs[bos:n - eos] = torch.softmax(logits[bos:n - eos], dim=0)
    return probs + decay * prev


# --------------------------------------------------------------------------
# main check
# --------------------------------------------------------------------------
def run(seq_len, steps, group_size, head_dim, topk_val, bos, eos, seed, decay=None):
    torch.manual_seed(seed)
    num_kv_heads = 1
    max_reqs = 2                      # two slots so page reuse can be exercised
    max_pages = seq_len + 8

    flow = (H2OTokenSparseAttention() if decay is None
            else H2OTokenSparseAttention(decay=decay))
    decay = flow.add.beta
    meta = flow.get_cache_meta_info(page_size=1, head_dim=head_dim)
    scale = flow.softmax.scale

    ctx = build_ctx(max_reqs, max_pages, num_kv_heads, group_size, head_dim,
                    topk_val, bos, eos)
    cctx = build_cache_ctx(num_kv_heads, head_dim,
                           ctx.max_num_pages + num_kv_heads)

    # paged KV pool, exactly as VTXGraphCachePool allocates it
    n_pages = ctx.max_num_pages + num_kv_heads          # + guard page
    cache_raw = {
        name: torch.zeros((n_pages, s[0], s[1]), dtype=DT, device=DEV)
        for name, s in meta.items()
    }
    wrap = lambda d: {k: as_vtensor(v, FORMAT.PAGED) for k, v in d.items()}

    # request 0 occupies token slots [0, seq_len); request 1 reuses slot range
    req_to_token = torch.zeros((max_reqs, max_pages), dtype=torch.int32, device=DEV)
    req_to_token[0, :seq_len] = torch.arange(seq_len, dtype=torch.int32, device=DEV)
    req_to_token[1, :seq_len] = torch.arange(seq_len, dtype=torch.int32, device=DEV)

    # ---------------- profile once, as the runtime does ----------------
    ctx.mode = Mode.profile
    q_dummy = as_vtensor(torch.empty((1, group_size, head_dim), device=DEV, dtype=DT),
                         FORMAT.BATCHED)
    o_dummy = as_vtensor(torch.empty((0, 1, 1), device=DEV, dtype=DT), FORMAT.RAGGED)
    dummy = {k: as_vtensor(torch.zeros((0, s[0], s[1]), dtype=DT, device=DEV), FORMAT.PAGED)
             for k, s in ({**meta, "k": (1, head_dim)}).items()}
    flow.forward_indexer(q_dummy, o_dummy, dummy, ctx=ctx)
    flow.forward_cache(cache=dummy, loc=torch.empty((0,), dtype=torch.int64, device=DEV),
                       ctx=cctx)
    cctx.mode = Mode.execute

    ctx.mode = Mode.execute
    results = {}

    # ---------------- fill the K cache for request 0 ----------------
    slots = torch.arange(seq_len, device=DEV)
    pages = page_ids(slots, num_kv_heads, 0)

    # A stable heavy-hitter set: keys carry a per-token affinity along a shared
    # direction that the queries also lean towards, so the same tokens attract
    # attention across steps. With i.i.d. keys and queries every token is
    # equally attended and the top-k boundary falls below one bf16 ulp, which
    # makes any top-k comparison meaningless regardless of implementation.
    u = torch.randn(head_dim, device=DEV)
    u = u / u.norm()
    affinity = torch.randn(seq_len, device=DEV) * 2.0
    k = (torch.randn(seq_len, head_dim, device=DEV)
         + affinity[:, None] * u[None, :]).to(DT)
    cache_raw["k"][pages, 0, :] = k

    # forward_cache runs once per written page -> hh_score starts at zero
    flow.forward_cache(cache=wrap(cache_raw),
                       loc=slots.to(torch.int64), ctx=cctx)

    seq_lens = torch.tensor([seq_len], dtype=torch.int32, device=DEV)
    # req_indices is int64 in the kernel contract; everything else int32
    plan_decode(seq_lens, req_to_token,
                torch.tensor([0], dtype=torch.int64, device=DEV), ctx)

    o = torch.zeros((ctx.max_num_pages, 1, 1), dtype=torch.int32, device=DEV)
    ref = torch.zeros(seq_len, dtype=torch.float64, device=DEV)
    traj = []

    for step in range(steps):
        q = (torch.randn(1, group_size, head_dim, device=DEV) * 0.3
             + u[None, None, :]).to(DT)
        flow.forward_indexer(as_vtensor(q, FORMAT.BATCHED),
                             as_vtensor(o, FORMAT.RAGGED),
                             wrap(cache_raw), ctx=ctx)
        torch.cuda.synchronize()
        ref = reference_step(ref, q.mean(dim=1).squeeze(0), k, scale, decay, bos, eos)
        if step in (0, steps // 4, steps // 2, steps - 1):
            got = cache_raw["hh_score"][pages, 0, 0].double()
            traj.append((step + 1, got.clone(), ref.clone()))

    results["mass"] = float(cache_raw["hh_score"][pages, 0, 0].double().sum())
    results["decay"] = decay
    results["scale"] = scale
    results["traj"] = traj

    # ---------------- page reuse: request 1 takes the same slots ----------------
    before = cache_raw["hh_score"][pages, 0, 0].clone()
    k2 = torch.randn(seq_len, head_dim, device=DEV, dtype=DT)
    cache_raw["k"][pages, 0, :] = k2
    flow.forward_cache(cache=wrap(cache_raw), loc=slots.to(torch.int64), ctx=cctx)
    torch.cuda.synchronize()
    after = cache_raw["hh_score"][pages, 0, 0].clone()
    results["reuse"] = (before, after)
    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seq-len", type=int, default=128)
    ap.add_argument("--steps", type=int, default=600)
    ap.add_argument("--group-size", type=int, default=4)
    ap.add_argument("--head-dim", type=int, default=128)
    ap.add_argument("--topk", type=int, default=32)
    ap.add_argument("--bos", type=int, default=4)
    ap.add_argument("--eos", type=int, default=8)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--decay", type=float, default=None,
                    help="override the flow decay (1.0 = undecayed sum)")
    a = ap.parse_args()

    print(f"device: {torch.cuda.get_device_name(0)}")
    r = run(a.seq_len, a.steps, a.group_size, a.head_dim, a.topk, a.bos, a.eos,
            a.seed, a.decay)
    print(f"flow: decay={r['decay']}  softmax scale={r['scale']:.6f} "
          f"(head_dim**-0.5 = {a.head_dim ** -0.5:.6f})\n")

    fails = []

    # steady state of sum_t decay^t * 1.0
    decayed = r["decay"] < 1.0
    expected_mass = 1.0 / (1.0 - r["decay"]) if decayed else float("inf")
    print(f"total probability mass in hh_score: {r['mass']:.2f} "
          + (f"(steady state 1/(1-decay) = {expected_mass:.2f})" if decayed
             else "(undecayed: grows by ~1.0 per step, without bound)"))
    # Each step injects exactly 1.0 of probability mass, so the recurrence
    # predicts the retained mass exactly. Falling short of it means increments
    # were rounded away -- which is how the undecayed accumulator fails: it
    # stops growing not because it is bounded but because it has frozen.
    predicted = expected_mass if decayed else float(a.steps)
    lost = 1.0 - r["mass"] / predicted
    print(f"  mass predicted by the recurrence = {predicted:.2f} "
          f"-> {lost:+.1%} lost to bf16 rounding")
    if lost > 0.10:
        fails.append(f"{lost:.0%} of the probability mass lost to rounding "
                     f"(retained {r['mass']:.1f} of {predicted:.1f})")
    print()
    print("hh_score vs an independent float64 reference of the same recurrence")
    print(f"{'step':>6} {'rel.err':>10} {'max |hh|':>10} {'top-k vs ref':>14} "
          f"{'boundary/ulp':>13}")
    for step, got, ref in r["traj"]:
        rel = float((got - ref).abs().max() / ref.abs().max())
        k = min(32, int((ref > 0).sum()))
        ov = len(set(got.topk(k).indices.tolist()) & set(ref.topk(k).indices.tolist())) / k
        srt = ref[ref > 0].sort(descending=True).values
        gap = float(srt[k - 1] - srt[k]) if srt.numel() > k else float("inf")
        ulp = float(srt[k - 1]) * 2 ** -8
        print(f"{step:>6} {rel:>10.2e} {float(got.max()):>10.2f} {ov:>13.0%} "
              f"{gap / ulp:>13.1f}")
        if step == r["traj"][-1][0]:
            if rel >= 0.15:
                fails.append(f"final rel.err {rel:.2e} >= 0.15")
            # only meaningful when the boundary is wider than bf16 can resolve
            if gap > ulp and ov < 1.0:
                fails.append(f"final top-{k} overlap {ov:.0%} < 100% "
                             f"despite a {gap/ulp:.1f}-ulp boundary")


    before, after = r["reuse"]
    reset_ok = bool((after == 0).all())
    print(f"\npage reuse: forward_cache on recycled pages")
    print(f"  hh_score before reset: max={float(before.max()):.3f} "
          f"nonzero={int((before != 0).sum())}/{before.numel()}")
    print(f"  hh_score after  reset: max={float(after.max()):.3f} "
          f"nonzero={int((after != 0).sum())}/{after.numel()}")
    if not reset_ok:
        fails.append("hh_score not zeroed on page reuse")
    if not bool((before != 0).any()):
        fails.append("nothing accumulated before reuse -- test is vacuous")

    print()
    for label, ok in (
        ("accumulates softmax probabilities (matches fp64 reference)",
         not any("rel.err" in f or "overlap" in f for f in fails)),
        ("accumulator retains its probability mass under bf16",
         not any("rounding" in f for f in fails)),
        ("forward_cache resets hh_score on page reuse",
         not any("zeroed" in f or "vacuous" in f for f in fails)),
    ):
        print(f"  [{'PASS' if ok else 'FAIL'}] {label}")

    print(f"\n{'ALL CHECKS PASSED' if not fails else 'FAILED: ' + '; '.join(fails)}")
    raise SystemExit(1 if fails else 0)


if __name__ == "__main__":
    main()

Happy to add this to the tree if you want it somewhere — it needs only the built extension and one GPU, so it would fit as a standalone check.