[Feature] WISE working-set decode for recurrent-depth models

benchmarks/wise_kernel_bench.py

@@ -0,0 +1,159 @@
+"""Decode-attention kernel timing for WISE modes (CUDA events, one layer call).
+
+Ouro-1.4B shapes (16 heads, head_dim 128, bf16). Sparse rows use a fixed
+support with the given density, newest block always kept. FlashAttention
+sparse rows attend the compacted page table of the same support (head union,
+identical across heads here, so the union equals each head's support).
+"""
+
+import argparse
+import json
+import math
+
+import torch
+
+from vllm_rlt.kernels.triton_attention import paged_attention
+from vllm_rlt.kernels.wise_attention import (
+    MODE_DENSE,
+    MODE_DISCOVER,
+    MODE_SPARSE,
+    wise_paged_attention,
+    wise_scratch,
+)
+from vllm_rlt.kernels.wise_flash import compact_page_tables
+
+
+def timeit(fn, iters):
+    """GPU time per call from CUDA-graph replays (no Python launch overhead, as in the engine)."""
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            fn()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        for _ in range(10):
+            fn()
+    for _ in range(2):
+        graph.replay()
+    torch.cuda.synchronize()
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    replays = max(1, iters // 10)
+    start.record()
+    for _ in range(replays):
+        graph.replay()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / (replays * 10) * 1000
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--cases", default="64x1024,40x2048,20x4096")
+    p.add_argument("--densities", default="0.17,0.33")
+    p.add_argument("--iters", type=int, default=50)
+    a = p.parse_args()
+    heads, dim = 16, 128
+    for case in a.cases.split(","):
+        rows, length = map(int, case.split("x"))
+        for page, backends in ((16, ("triton", 3, 4)), (256, (2,))):
+            pages = math.ceil(length / page)
+            k = torch.randn((rows * pages, page, heads, dim), device="cuda", dtype=torch.bfloat16)
+            v = torch.randn_like(k)
+            q = torch.randn((rows, heads, dim), device="cuda", dtype=torch.bfloat16)
+            tables = (
+                torch.randperm(rows * pages, device="cuda")
+                .to(torch.int32)
+                .view(rows, pages)
+                .contiguous()
+            )
+            lengths = torch.full((rows,), length, device="cuda", dtype=torch.int32)
+            max_blocks = 1 << math.ceil(math.log2(math.ceil(pages * page / 32)))
+            slots = torch.arange(rows, device="cuda", dtype=torch.int32)
+            scratch = wise_scratch(rows, heads, dim, max_blocks, "cuda")
+            stats = torch.zeros(2, device="cuda", dtype=torch.int64)
+            out = {"rows": rows, "length": length, "page": page}
+            for backend in backends:
+                if backend == "triton":
+                    out["triton_dense"] = timeit(
+                        lambda: paged_attention(q, k, v, tables, lengths), a.iters
+                    )
+                    for mode, name in (
+                        (MODE_DENSE, "wise_split_dense"),
+                        (MODE_DISCOVER, "wise_discover"),
+                    ):
+                        modes = torch.full((rows,), mode, device="cuda", dtype=torch.int32)
+                        support = torch.ones(
+                            (rows, 1, heads, max_blocks), device="cuda", dtype=torch.uint8
+                        )
+                        out[name] = timeit(
+                            lambda: wise_paged_attention(
+                                q, k, v, tables, lengths, modes, slots, support, 0, scratch, stats
+                            ),
+                            a.iters,
+                        )
+                from vllm_rlt.kernels.flash_attention import FlashPagedAttention
+
+                flash = None
+                if backend != "triton":
+                    try:
+                        flash = FlashPagedAttention(
+                            torch.device("cuda"), torch.bfloat16, dim, page, f"flash_attn_{backend}"
+                        )
+                    except (ImportError, ValueError) as error:
+                        out[f"fa{backend}"] = str(error)
+                        continue
+                    out[f"fa{backend}_dense"] = timeit(
+                        lambda: flash(q, k, v, tables, lengths), a.iters
+                    )
+                    out[f"fa{backend}_dense_lse"] = timeit(
+                        lambda: flash(q, k, v, tables, lengths, return_lse=True), a.iters
+                    )
+                for density in map(float, a.densities.split(",")):
+                    blocks = torch.arange(max_blocks, device="cuda")
+                    count = math.ceil(length / 32)
+                    stride = max(1, round(1 / density))
+                    chosen = ((blocks % stride == 0) | (blocks == count - 1)) & (blocks < count)
+                    support = (
+                        chosen[None, None, None, :]
+                        .expand(rows, 1, heads, -1)
+                        .to(torch.uint8)
+                        .contiguous()
+                    )
+                    modes = torch.full((rows,), MODE_SPARSE, device="cuda", dtype=torch.int32)
+                    if backend == "triton":
+                        out[f"wise_sparse_{density}"] = timeit(
+                            lambda: wise_paged_attention(
+                                q, k, v, tables, lengths, modes, slots, support, 0, scratch, stats
+                            ),
+                            a.iters,
+                        )
+                    else:
+                        ot, ol = torch.empty_like(tables), torch.empty_like(lengths)
+
+                        def run():
+                            t, n = compact_page_tables(
+                                tables, lengths, modes, slots, support, 0, page, ot, ol
+                            )
+                            return flash(q, k, v, t, n)
+
+                        out[f"fa{backend}_sparse_{density}"] = timeit(run, a.iters)
+                        out[f"fa{backend}_compact_only_{density}"] = timeit(
+                            lambda: compact_page_tables(
+                                tables, lengths, modes, slots, support, 0, page, ot, ol
+                            ),
+                            a.iters,
+                        )
+            print(
+                json.dumps(
+                    {
+                        key: round(val, 1) if isinstance(val, float) else val
+                        for key, val in out.items()
+                    }
+                ),
+                flush=True,
+            )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/wise_quality.py

@@ -0,0 +1,138 @@
+"""Answer quality of wise-cached-decode through the engine, paired with dense decode.
+
+Runs the fixed 200 Ouro plain-completion HotpotQA/2Wiki prompts greedily for 32
+tokens in both arms (same engine settings, Triton, always 4 loops) and scores
+the first line with the WISE authors' benchmark_scores (needs wise-reference on
+PYTHONPATH). Reports per-row predictions, F1/EM and token agreement.
+"""
+
+import argparse
+import json
+
+import torch
+from transformers import AutoTokenizer
+from wise.metrics import benchmark_scores
+
+from vllm_rlt import LLM, SamplingParams
+from vllm_rlt.config import CacheConfig, ExecutionConfig, ExitConfig, SchedulerConfig
+from vllm_rlt.engine.stage_engine import StageEngine
+from vllm_rlt.models import OuroForCausalLM
+from vllm_rlt.wise import WiseConfig
+
+
+class InProcessLLM(LLM):
+    """LLM facade over an in-process StageEngine, so the script can read engine state."""
+
+    def __init__(self, model, **options):
+        self.tokenizer = None
+        self._tokenizer_source = None
+        self._next_request_id = 0
+        self.engine = StageEngine(model, **options)
+
+
+def run_arm(a, rows, arm):
+    model = OuroForCausalLM.from_pretrained(a.model, device="cuda", dtype=torch.bfloat16)
+    # EOS stops generation; scoring cuts at <|im_end|> too and keeps the first line.
+    params = SamplingParams(max_tokens=a.max_tokens, exit_threshold=1.0, min_loops=4, max_loops=4)
+    longest = max(len(r["token_ids"]) for r in rows) + a.max_tokens + 32
+    seqs = 32
+    # Both arms: the benchmark's best configuration (async scheduling + CUDA graphs).
+    llm = InProcessLLM(
+        model,
+        attention_backend=a.backend,
+        cache_config=CacheConfig(num_blocks=seqs * ((longest + 15) // 16) * 4, block_size=16),
+        scheduler_config=SchedulerConfig(
+            max_num_seqs=seqs, max_num_batched_tokens=4096, prefill_chunk_size=2048
+        ),
+        execution_config=ExecutionConfig(async_scheduling=True, cuda_graphs=True),
+        exit_config=ExitConfig(mode="ouro_delayed"),
+        wise_config=WiseConfig(max_model_len=longest, max_slots=seqs + 8, kernel=a.wise_kernel)
+        if arm == "wise"
+        else None,
+    )
+    outputs = llm.generate([r["token_ids"] for r in rows], params)
+    wise = llm.engine.cache_manager.wise
+    return [o.token_ids for o in outputs], wise.mean_density() if wise is not None else None
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", required=True)
+    p.add_argument("--prompts", required=True)
+    p.add_argument("--out", required=True)
+    p.add_argument("--max-tokens", type=int, default=32)
+    p.add_argument("--backend", default="flash_attn_4", help="prefill and dense decode backend")
+    p.add_argument("--wise-kernel", default="triton", choices=("auto", "triton"))
+    p.add_argument("--arm", choices=("dense", "wise"), help=argparse.SUPPRESS)
+    a = p.parse_args()
+    rows = [json.loads(x) for x in open(a.prompts)]
+    if a.arm:
+        tokens, density = run_arm(a, rows, a.arm)
+        with open(a.out, "x") as f:
+            json.dump({"tokens": tokens, "density": density}, f)
+        return
+    import subprocess
+    import sys
+    import tempfile
+
+    tokenizer = AutoTokenizer.from_pretrained(a.model)
+    results = {}
+    with tempfile.TemporaryDirectory() as tmp:
+        for arm in ("dense", "wise"):
+            path = f"{tmp}/{arm}.json"
+            subprocess.run(
+                [
+                    sys.executable,
+                    __file__,
+                    "--model",
+                    a.model,
+                    "--prompts",
+                    a.prompts,
+                    "--out",
+                    path,
+                    "--max-tokens",
+                    str(a.max_tokens),
+                    "--backend",
+                    a.backend,
+                    "--wise-kernel",
+                    a.wise_kernel,
+                    "--arm",
+                    arm,
+                ],
+                check=True,
+            )
+            results[arm] = json.load(open(path))
+    with open(a.out, "x") as f:
+        for row, dense, wise in zip(rows, results["dense"]["tokens"], results["wise"]["tokens"]):
+            record = {
+                "id": row["id"],
+                "benchmark": row["benchmark"],
+                "format": row.get("format"),
+                "identical_tokens": dense == wise,
+                "first_divergence": next(
+                    (i for i, (u, v) in enumerate(zip(dense, wise)) if u != v),
+                    min(len(dense), len(wise)),
+                ),
+            }
+            end = tokenizer.convert_tokens_to_ids("<|im_end|>")
+            for arm, tokens in (("dense", dense), ("wise", wise)):
+                tokens = tokens[: tokens.index(end)] if end in tokens else tokens
+                text = tokenizer.decode(tokens, skip_special_tokens=True).split("\n", 1)[0].strip()
+                f1, em = benchmark_scores(row["benchmark"], text, row["gold_answer"])
+                record[arm] = {"prediction": text, "f1": f1, "em": em, "tokens": tokens}
+            f.write(json.dumps(record) + "\n")
+        f.write(
+            json.dumps(
+                {
+                    "summary": True,
+                    "wise_mean_support_density": results["wise"]["density"],
+                    "backend": a.backend,
+                    "wise_kernel": a.wise_kernel,
+                }
+            )
+            + "\n"
+        )
+
+
+if __name__ == "__main__":
+    main()