diff --git a/benchmarks/wise_kernel_bench.py b/benchmarks/wise_kernel_bench.py new file mode 100644 index 00000000..f540e7e5 --- /dev/null +++ b/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() diff --git a/benchmarks/wise_quality.py b/benchmarks/wise_quality.py new file mode 100644 index 00000000..096fa43e --- /dev/null +++ b/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() diff --git a/benchmarks/wise_throughput.py b/benchmarks/wise_throughput.py new file mode 100644 index 00000000..094d2550 --- /dev/null +++ b/benchmarks/wise_throughput.py @@ -0,0 +1,304 @@ +"""High-concurrency decode throughput of Ouro with and without WISE (wise-cached-decode). + +Each case runs C requests of exactly L real prompt tokens (Ouro-tokenized +HotpotQA/2Wiki plain prompts, concatenated and left-truncated to L) with greedy +decoding, ignore_eos and a fixed number of output tokens, always 4 loops. +Decode throughput = (tokens generated after t_first) / (t_done - t_first), where +t_first is the first engine step after which every request has its first token +(prefill complete) and t_done is when all have finished. Tokens that early requests +generate while later prompts still prefill fall before t_first and are not counted. +Arms run in fresh processes, interleaved per repeat. + +Arms (all Triton unless noted; graph arms use CUDA graphs up to the case batch) + dense_eager synchronous eager decode + dense_graph synchronous decode with CUDA graphs + dense_best async scheduling + CUDA graphs + dense_fa3 FlashAttention-3 with async scheduling + CUDA graphs + wise_eager WISE decode, eager (ouro4: discover at depths 0-1, reuse at 2-3) + wise WISE decode with CUDA graphs + wise_best WISE decode with async scheduling + CUDA graphs + wise_all WISE discovery and selection with CUDA graphs, every block kept + split_dense CUDA graphs, every row dense through WISE's split-KV kernel (its dense baseline) + fixed33 CUDA graphs, no discovery, fixed one-in-three block support at reuse depths + spec_trunc_kN CUDA graphs, self-speculation: N draft tokens read out after 2 loops + spec_wise_kN CUDA graphs, self-speculation: N draft tokens from WISE loops 3-4; both + speculative arms verify densely, so greedy output equals dense decode +Any arm may end in @fa2, @fa3 or @fa4 to use FlashAttention 2/3/4 instead of Triton +(FA2 pages are 256 tokens, others 16); dense_fa3 is dense_best@fa3. A WISE arm with a +"_tk" suffix (e.g. wise_best_tk@fa3) keeps FlashAttention for prefill but decodes +WISE with the Triton split-KV kernel. +""" + +import argparse +import json +import math +import os +import time + +import torch + +from vllm_rlt import LLM, SamplingParams +from vllm_rlt.config import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SchedulerConfig, + SpeculativeConfig, +) +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.wise import WiseConfig + +ARMS = ( + "dense_eager", + "dense_graph", + "dense_best", + "dense_fa3", + "wise_eager", + "wise", + "wise_all", + "fixed33", + "spec_trunc_k2", + "spec_trunc_k4", + "spec_wise_k2", + "spec_wise_k4", +) + + +def prompts_of_length(rows, length, count): + out = [] + for i in range(count): + tokens, j = [], i + while len(tokens) < length: + tokens = rows[j % len(rows)]["token_ids"] + tokens + j += 1 + out.append(tokens[-length:]) + return out + + +class InProcessLLM(LLM): + """LLM facade over an in-process StageEngine, so the benchmark 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 close(self): + self.engine.close() + + +def build(model, arm, concurrency, length, output_tokens): + wise = None + if arm == "dense_fa3": + arm = "dense_best@fa3" + arm, _, tag = arm.partition("@") + triton_kernel = arm.endswith("_tk") + arm = arm.removesuffix("_tk") + backend = {"": "triton", "fa2": "flash_attn_2", "fa3": "flash_attn_3", "fa4": "flash_attn_4"}[ + tag + ] + page = 256 if tag == "fa2" else 16 + graphs = arm not in ("dense_eager", "wise_eager") + asynchronous = arm in ("dense_best", "wise_best") + execution = ExecutionConfig( + async_scheduling=asynchronous, + cuda_graphs=graphs, + cuda_graph_max_batch_size=max(128, concurrency), + ) + speculative = None + if arm.startswith("spec_"): + k = int(arm.rsplit("_k", 1)[1]) + draft = "wise" if arm.startswith("spec_wise") else "truncate" + speculative = SpeculativeConfig(k, draft_attention=draft) + if draft == "wise": + wise = WiseConfig( + max_model_len=length + output_tokens + 32 + k, max_slots=concurrency + 8 + ) + elif arm.startswith("wise") or arm in ("fixed33", "split_dense"): + override = { + "wise": None, + "wise_best": None, + "wise_eager": None, + "wise_all": "all", + "fixed33": 1 / 3, + "split_dense": "dense", + }[arm] + wise = WiseConfig( + max_model_len=length + output_tokens + 32, + max_slots=concurrency + 8, + support_override=override, + kernel="triton" if triton_kernel else "auto", + ) + return InProcessLLM( + model, + attention_backend=backend, + # Identical, case-sized KV pool for every arm: all C requests fit at once. + cache_config=CacheConfig( + num_blocks=math.ceil( + concurrency * math.ceil((length + output_tokens + 32) / page) * 4 * 1.05 + ), + block_size=page, + ), + scheduler_config=SchedulerConfig( + max_num_seqs=concurrency, max_num_batched_tokens=4096, prefill_chunk_size=2048 + ), + execution_config=execution, + # Async needs a delayed exit policy and speculation the fixed-depth ouro policy; + # threshold 1.0 never exits, so every arm runs 4 loops. + exit_config=ExitConfig(mode="ouro" if speculative else "ouro_delayed"), + speculative_config=speculative, + wise_config=wise, + ) + + +def timed(llm, prompts, tokens, tag="r"): + params = SamplingParams( + max_tokens=tokens, ignore_eos=True, exit_threshold=1.0, min_loops=4, max_loops=4 + ) + engine = llm.engine + ids = [f"{tag}{i}" for i in range(len(prompts))] + torch.cuda.synchronize() + start = time.perf_counter() + for request_id, prompt in zip(ids, prompts): + engine.add_request(request_id, list(prompt), params) + first, outputs = None, {} + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + outputs[output.request_id] = output.token_ids + if first is None: + live = engine.scheduler.requests.values() + if all(len(r.generated_token_ids) >= 1 for r in live): + torch.cuda.synchronize() + first = time.perf_counter() + # Early requests may decode while later prompts still prefill; those + # tokens belong to the prefill interval, not to the decode count. + requests = engine.scheduler.requests + before = sum( + len(outputs[i]) if i in outputs else len(requests[i].generated_token_ids) + for i in ids + ) + torch.cuda.synchronize() + done = time.perf_counter() + tokens_out = [outputs[i] for i in ids] + assert all(len(t) == tokens for t in tokens_out) + return first - start, done - first, tokens_out, len(ids) * tokens - before + + +def run_one(a): + """One (case, arm, repeat) in a fresh process, so every arm starts from empty GPU memory.""" + rows = [json.loads(x) for x in open(a.prompts)] + length, concurrency = map(int, a.cases.split("x")) + prompts = prompts_of_length(rows, length, concurrency) + model = OuroForCausalLM.from_pretrained(a.model, device="cuda", dtype=torch.bfloat16) + llm = build(model, a.arms, concurrency, length, a.output_tokens) + blocks = llm.engine.cache_manager.num_free_blocks + # Same batch size, short prompts: compiles kernels and captures decode graphs untimed. + timed(llm, [p[-64:] for p in prompts], 8, tag="warm") + wise = llm.engine.cache_manager.wise + if wise is not None: + wise.stats.zero_() # support density of the timed cohort only + prefill, decode, tokens, decoded = timed(llm, prompts, a.output_tokens) + record = { + "length": length, + "concurrency": concurrency, + "arm": a.arms, + "repeat": a.repeats, + "output_tokens": a.output_tokens, + "prefill_s": prefill, + "decode_s": decode, + "decode_tokens": decoded, + "decode_tok_s": decoded / decode, + "end_to_end_tok_s": concurrency * a.output_tokens / (prefill + decode), + "free_blocks_at_start": blocks, + "support_density": wise.mean_density() if wise is not None else None, + "attention": llm.engine.cache_manager.attention_info.get("selected_backend"), + "speculative": ( + vars(llm.engine.speculative_runner.stats).copy() + if llm.engine.speculative_runner is not None + else None + ), + "tokens": tokens, + } + with open(a.out, "x") as f: + json.dump(record, f) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", required=True) + p.add_argument("--prompts", required=True, help="JSONL with token_ids (Ouro plain prompts)") + p.add_argument( + "--cases", required=True, help="comma list LENGTHxCONCURRENCY, e.g. 1024x32,2048x16" + ) + p.add_argument("--arms", default=",".join(ARMS)) + p.add_argument("--output-tokens", type=int, default=256) + p.add_argument("--repeats", type=int, default=2) + p.add_argument("--out", required=True) + p.add_argument("--single", action="store_true", help=argparse.SUPPRESS) + a = p.parse_args() + if a.single: + return run_one(a) + import subprocess + import sys + import tempfile + + arms = a.arms.split(",") + with open(a.out, "x") as f, tempfile.TemporaryDirectory() as tmp: + for case in a.cases.split(","): + for repeat in range(1, a.repeats + 1): # each process warms up before timing + order = arms if repeat % 2 else arms[::-1] + records = {} + for arm in order: + path = os.path.join(tmp, f"{case}-{arm}-{repeat}.json") + command = [ + sys.executable, + __file__, + "--single", + "--model", + a.model, + "--prompts", + a.prompts, + "--cases", + case, + "--arms", + arm, + "--output-tokens", + str(a.output_tokens), + "--repeats", + str(repeat), + "--out", + path, + ] + done = subprocess.run(command, capture_output=True, text=True) + if done.returncode: + records[arm] = { + "case": case, + "arm": arm, + "repeat": repeat, + "error": done.stderr[-2000:], + } + else: + records[arm] = json.load(open(path)) + base = records.get("dense_eager", {}).get("tokens") + for arm in order: + record = records[arm] + mine = record.pop("tokens", None) + if base is not None and mine is not None: + record["sequences_identical_to_dense_eager"] = sum( + x == y for x, y in zip(mine, base) + ) + record["first_divergence_mean"] = sum( + next((i for i, (u, v) in enumerate(zip(x, y)) if u != v), len(x)) + for x, y in zip(mine, base) + ) / len(base) + f.write(json.dumps(record) + "\n") + f.flush() + os.fsync(f.fileno()) + print(json.dumps(record)[:400], flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/test_speculative.py b/tests/test_speculative.py index b2593c83..0807a9b3 100644 --- a/tests/test_speculative.py +++ b/tests/test_speculative.py @@ -484,6 +484,13 @@ def test_invalid_k(k): SpeculativeConfig(k) +def test_wise_draft_requires_block_speculation(): + # Pipeline rounds have no WISE draft path; only block speculation accepts it. + assert SpeculativeConfig(4, draft_attention="wise").mode == "block" + with pytest.raises(ValueError, match="requires block speculation"): + SpeculativeConfig(4, draft_loops=1, mode="pipeline", draft_attention="wise") + + @pytest.mark.gpu @pytest.mark.parametrize("backend", ["triton", "flash_attn_4"]) def test_gpu_sampling_and_greedy_match_replay_with_ragged_requests(backend): diff --git a/tests/test_wise.py b/tests/test_wise.py new file mode 100644 index 00000000..25aaebcb --- /dev/null +++ b/tests/test_wise.py @@ -0,0 +1,548 @@ +import math + +import pytest +import torch +import torch.nn.functional as F + +# The WISE kernels are Triton; CPU-only installs (CI) skip this module. +pytest.importorskip("triton") + +from tests.helpers import LocalLLM as LLM +from tests.helpers import tiny_ouro_config +from vllm_rlt import SamplingParams +from vllm_rlt.config import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SchedulerConfig, + SpeculativeConfig, +) +from vllm_rlt.kernels.triton_attention import paged_attention +from vllm_rlt.kernels.wise_attention import ( + MODE_DENSE, + MODE_DISCOVER, + MODE_DISCOVER_UNION, + MODE_SPARSE, + block_mass, + select_support, + wise_paged_attention, + wise_scratch, +) +from vllm_rlt.models import OuroForCausalLM +from vllm_rlt.request import RequestPhase +from vllm_rlt.wise import WiseConfig + +pytestmark = pytest.mark.gpu + + +def _pool(rows=6, heads=4, kv_heads=2, dim=64, page=16, length=400, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + pages = math.ceil(length / page) + k = torch.randn( + (rows * pages + 3, page, kv_heads, dim), device="cuda", generator=g, dtype=torch.bfloat16 + ) + v = torch.randn_like(k) + q = torch.randn((rows, heads, dim), device="cuda", generator=g, dtype=torch.bfloat16) + perm = torch.randperm(rows * pages, device="cuda", generator=g).to(torch.int32) + 3 + tables = perm.view(rows, pages).contiguous() + lengths = torch.tensor([length - 37 * r for r in range(rows)], device="cuda", dtype=torch.int32) + return q, k, v, tables, lengths + + +def _dense_probs(q, k, v, tables, lengths, row): + n = int(lengths[row]) + pos = torch.arange(n, device="cuda") + page = k.shape[1] + blocks = tables[row, pos // page].long() + keys = k[blocks, pos % page].float().repeat_interleave(q.shape[1] // k.shape[2], 1) + values = v[blocks, pos % page].float().repeat_interleave(q.shape[1] // k.shape[2], 1) + scores = torch.einsum("hd,thd->ht", q[row].float(), keys) / math.sqrt(q.shape[-1]) + return scores, values + + +def _support_reference(probabilities, block_size=32, eta=0.95): + """vime_plugins/ouro/wise_ouro_reference._support for one query row (heads, 1, n).""" + heads, n, _ = probabilities.shape + blocks = math.ceil(n / block_size) + padded = F.pad(probabilities.float(), (0, blocks * block_size - n, 0, blocks * block_size - n)) + mass = padded.reshape(heads, blocks, block_size, blocks, block_size).sum(-1).sum(2) + counts = torch.full((blocks,), block_size, device=mass.device, dtype=mass.dtype) + counts[-1] = n - (blocks - 1) * block_size + mass = mass / counts[None, :, None] + causal = torch.ones(blocks, blocks, device=mass.device, dtype=torch.bool).tril() + ranked, order = torch.sort(mass.masked_fill(~causal, -torch.inf), descending=True, stable=True) + keep = ranked.clamp_min(0).cumsum(-1).lt((eta * mass.sum(-1)).unsqueeze(-1)).sum(-1) + 1 + keep = torch.minimum(keep, torch.arange(1, blocks + 1, device=mass.device)).clamp_min(1) + selected = torch.arange(blocks, device=mass.device) < keep.unsqueeze(-1) + return ( + torch.zeros_like(causal).expand(heads, -1, -1).clone().scatter_(-1, order, selected) + & causal + ) + + +def _launch(q, k, v, tables, lengths, mode, support=None, splits=None, num_warps=1): + rows, heads = q.shape[:2] + blocks = 1 << math.ceil(math.log2(math.ceil(tables.shape[1] * k.shape[1] / 32))) + if support is None: + support = torch.ones((rows, 1, heads, blocks), device="cuda", dtype=torch.uint8) + modes = torch.full((rows,), mode, device="cuda", dtype=torch.int32) + slots = torch.arange(rows, device="cuda", dtype=torch.int32) + scratch = wise_scratch(rows, heads, q.shape[-1], support.shape[-1], "cuda") + stats = torch.zeros(2, device="cuda", dtype=torch.int64) + out = wise_paged_attention( + q, + k, + v, + tables, + lengths, + modes, + slots, + support, + 0, + scratch, + stats, + splits=splits, + num_warps=num_warps, + ) + return out, scratch["mass"][0], scratch["mass"][1], support, stats + + +def test_split_kv_modes_agree_and_one_split_is_bitwise_the_dense_kernel(): + q, k, v, tables, lengths = _pool() + expected = paged_attention(q, k, v, tables, lengths) + # Same warps as the dense kernel: one split is the dense kernel bit for bit. + assert torch.equal( + _launch(q, k, v, tables, lengths, MODE_DENSE, splits=1, num_warps=4)[0], expected + ) + for splits in (None, 4): + dense = _launch(q, k, v, tables, lengths, MODE_DENSE, splits=splits)[0] + torch.testing.assert_close(dense, expected, atol=2e-2, rtol=2e-2) + for mode in (MODE_DISCOVER, MODE_DISCOVER_UNION, MODE_SPARSE): + assert torch.equal(_launch(q, k, v, tables, lengths, mode, splits=splits)[0], dense), ( + mode + ) + + +def test_fused_selection_matches_softmax_block_mass_and_reference_support(): + q, k, v, tables, lengths = _pool() + support = torch.full((q.shape[0], 1, q.shape[1], 16), 7, device="cuda", dtype=torch.uint8) + _, mm, ms, support, stats = _launch(q, k, v, tables, lengths, MODE_DISCOVER, support, splits=4) + mass, valid = block_mass(mm, ms, lengths.long()) + expected = select_support(mass, valid, 0.95) + assert torch.equal(support[:, 0].bool(), expected) + assert int(stats[0]) == int(expected.sum()) and int(stats[1]) == int(valid.sum()) * q.shape[1] + for row in range(q.shape[0]): + scores, _ = _dense_probs(q, k, v, tables, lengths, row) + probs = scores.softmax(-1) + n = probs.shape[-1] + blocks = math.ceil(n / 32) + sums = F.pad(probs, (0, blocks * 32 - n)).view(probs.shape[0], blocks, 32).sum(-1) + torch.testing.assert_close(mass[row, :, :blocks], sums, atol=2e-6, rtol=1e-5) + # Reference on a one-row query block: place the row last, pad earlier rows. + full = torch.zeros((probs.shape[0], n, n), device="cuda") + full[:, -1] = probs + reference = _support_reference(full)[:, -1] + assert torch.equal(expected[row, :, :blocks], reference) + + +def test_union_discovery_ors_into_the_first_support(): + q, k, v, tables, lengths = _pool(seed=1) + q2 = torch.randn_like(q) + _, mm, ms, first, _ = _launch( + q, + k, + v, + tables, + lengths, + MODE_DISCOVER, + torch.zeros((q.shape[0], 1, q.shape[1], 16), device="cuda", dtype=torch.uint8), + ) + a = select_support(*block_mass(mm, ms, lengths.long()), 0.95) + before = first.clone() + _, mm, ms, union, stats = _launch(q2, k, v, tables, lengths, MODE_DISCOVER_UNION, first) + b = select_support(*block_mass(mm, ms, lengths.long()), 0.95) + assert torch.equal(union[:, 0].bool(), a | b) + assert torch.equal(before[:, 0].bool(), a) + # Density counts the reused union: the union adds only the blocks it newly sets. + assert stats.tolist() == [int((a | b).sum() - a.sum()), 0] + + +def test_sparse_rows_equal_masked_softmax_over_selected_blocks(): + q, k, v, tables, lengths = _pool(seed=3) + rows, heads = q.shape[:2] + g = torch.Generator(device="cuda").manual_seed(7) + support = (torch.rand((rows, 1, heads, 16), device="cuda", generator=g) < 0.35).to(torch.uint8) + support[..., 0] = 1 + out = _launch(q, k, v, tables, lengths, MODE_SPARSE, support=support)[0] + for row in range(rows): + scores, values = _dense_probs(q, k, v, tables, lengths, row) + n = scores.shape[-1] + keep = support[row, 0].bool().repeat_interleave(32, -1)[:, :n] + weights = scores.masked_fill(~keep, -torch.inf).softmax(-1) + expected = torch.einsum("ht,thd->hd", weights, values) + torch.testing.assert_close(out[row].float(), expected, atol=2e-2, rtol=2e-2) + + +def _llm(model, wise, graphs=False): + return LLM( + model, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=512, prefill_chunk_size=512 + ), + execution_config=ExecutionConfig(cuda_graphs=graphs), + wise_config=wise, + ) + + +@pytest.fixture(scope="module") +def model(): + torch.manual_seed(0) + config = tiny_ouro_config( + head_dim=32, + num_attention_heads=4, + num_key_value_heads=4, + hidden_size=128, + max_position_embeddings=1024, + ) + return OuroForCausalLM(config).to("cuda", torch.bfloat16) + + +def test_engine_all_blocks_wise_matches_dense_engine_and_real_wise_prunes(model): + prompts = [list((torch.arange(300 + 50 * i) * (i + 3)) % 60 + 2) for i in range(3)] + prompts = [[int(t) for t in p] for p in prompts] + params = SamplingParams(max_tokens=12, ignore_eos=True) + dense = _llm( + model, WiseConfig(max_model_len=1024, max_slots=8, support_override="dense") + ).generate(prompts, params) + forced = _llm( + model, WiseConfig(max_model_len=1024, max_slots=8, support_override="all") + ).generate(prompts, params) + assert [o.token_ids for o in forced] == [o.token_ids for o in dense] + llm = _llm(model, WiseConfig(max_model_len=1024, max_slots=8)) + real = llm.generate(prompts, params) + assert all(len(o.token_ids) == 12 for o in real) + assert llm.engine.cache_manager.wise.slots == {} + + +def test_cuda_graph_wise_replays_eager_wise(model): + prompts = [[int(t) for t in (torch.arange(260 + 70 * i) * (i + 5)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=10, ignore_eos=True) + eager = _llm(model, WiseConfig(max_model_len=1024, max_slots=8)).generate(prompts, params) + llm = _llm(model, WiseConfig(max_model_len=1024, max_slots=8), graphs=True) + graphed = llm.generate(prompts, params) + assert [o.token_ids for o in graphed] == [o.token_ids for o in eager] + assert llm.engine.model_runner.graphs.replays > 0 + dense = _llm( + model, WiseConfig(max_model_len=1024, max_slots=8, support_override="dense"), graphs=True + ).generate(prompts, params) + forced = _llm( + model, WiseConfig(max_model_len=1024, max_slots=8, support_override="all"), graphs=True + ) + assert [o.token_ids for o in forced.generate(prompts, params)] == [o.token_ids for o in dense] + + +@pytest.mark.parametrize("graphs", [False, True]) +def test_wise_draft_speculation_is_lossless_under_greedy(model, graphs): + prompts = [[int(t) for t in (torch.arange(280 + 60 * i) * (i + 7)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=14, ignore_eos=True) + dense = _llm(model, None).generate(prompts, params) + llm = LLM( + model, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=512, prefill_chunk_size=512 + ), + execution_config=ExecutionConfig(cuda_graphs=graphs), + speculative_config=SpeculativeConfig(3, draft_attention="wise"), + wise_config=WiseConfig(max_model_len=1024, max_slots=8), + ) + drafted = llm.generate(prompts, params) + assert [o.token_ids for o in drafted] == [o.token_ids for o in dense] + stats = llm.engine.speculative_runner.stats + assert stats.drafted_tokens > 0 + assert llm.engine.cache_manager.wise.mean_density() < 1 + + +@pytest.mark.parametrize("graphs", [False, True]) +def test_async_wise_matches_synchronous_wise(model, graphs): + prompts = [[int(t) for t in (torch.arange(260 + 70 * i) * (i + 5)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=10, ignore_eos=True) + + def engine(asynchronous): + return LLM( + model, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=512, prefill_chunk_size=512 + ), + execution_config=ExecutionConfig(async_scheduling=asynchronous, cuda_graphs=graphs), + exit_config=ExitConfig(mode="ouro_delayed"), + wise_config=WiseConfig(max_model_len=1024, max_slots=8), + ) + + sync = [o.token_ids for o in engine(False).generate(prompts, params)] + llm = engine(True) + assert [o.token_ids for o in llm.generate(prompts, params)] == sync + assert llm.engine.cache_manager.wise.mean_density() < 1 + + +def test_incomplete_discovery_falls_back_to_dense(model): + llm = _llm(model, WiseConfig(max_model_len=1024, max_slots=8)) + wise = llm.engine.cache_manager.wise + + class Batch: + allocations = (("x", object()),) + + def plan(depth, position): + b = Batch() + b.rows = ((b.allocations[0][1], depth, position),) + b.context_lengths = torch.tensor([position + 1], device="cuda", dtype=torch.int32) + b.block_tables = torch.zeros((1, 64), device="cuda", dtype=torch.int32) + wise.mark(b, [True]) + return wise.plan(b)["modes"].tolist() + + # Depth 0 ran in a packed batch (no support recorded): depth 1 and the reuse + # depths of that token attend densely; the next token discovers again. + assert plan(1, 5) == [MODE_DENSE] and plan(2, 5) == [MODE_DENSE] + assert plan(0, 6) == [MODE_DISCOVER] and plan(1, 6) == [MODE_DISCOVER_UNION] + assert plan(2, 6) == [MODE_SPARSE] + # Support released between discovery and reuse (KV offload) or another position's. + wise.release("x") + assert plan(3, 6) == [MODE_DENSE] and plan(2, 7) == [MODE_DENSE] + assert wise.fallbacks == 4 + + +def test_offload_between_discovery_and_reuse_resumes_densely(model): + prompt = [int(t) for t in (torch.arange(300) * 3) % 60 + 2] + # Every token runs all four depths, so suspension can fall between 1 and 2. + params = SamplingParams( + max_tokens=8, ignore_eos=True, exit_threshold=1.0, min_loops=4, max_loops=4 + ) + + def engine(override): + return LLM( + model, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=1, + max_num_batched_tokens=512, + prefill_chunk_size=512, + enable_preemption=True, + ), + wise_config=WiseConfig(max_model_len=1024, max_slots=8, support_override=override), + ).engine + + dense = engine("dense") + dense.add_request("r", prompt, params) + while dense.has_unfinished_requests(): + outputs = [o for o in dense.step() if o.finished] + reference = outputs[0].token_ids + + e = engine("all") + e.add_request("r", prompt, params) + request = e.scheduler.requests["r"] + discovery = len(e.cache_manager.wise.config.discovery_depths) + for _ in range(200): + e.step() + if ( + request.phase == RequestPhase.RECURRENT + and request.loops_done == discovery + and request.generated_token_ids + ): + break + else: + pytest.fail("request did not stop between discovery and reuse") + e.add_request("pressure", [3, 4], SamplingParams(max_tokens=1, ignore_eos=True)) + assert e.scheduler.preempt(e.scheduler.requests["pressure"]) + assert e.preemption.snapshots["r"]["kv"] is not None + e.abort_request("pressure") + finished = [] + while e.has_unfinished_requests(): + finished += [o for o in e.step() if o.finished] + assert e.preemption.resumptions == 1 + assert e.cache_manager.wise.fallbacks > 0 + assert [o.token_ids for o in finished] == [reference] + + +def _flash(generation, page): + from vllm_rlt.kernels.flash_attention import FlashPagedAttention + + try: + return FlashPagedAttention( + torch.device("cuda"), torch.bfloat16, 64, page, f"flash_attn_{generation}" + ) + except (ImportError, ValueError) as error: + pytest.skip(str(error)) + + +FLASH = [(2, 256), (3, 16), (4, 16)] + + +@pytest.mark.parametrize("generation,page", FLASH) +def test_flash_all_blocks_compaction_is_bitwise_dense_flash(generation, page): + from vllm_rlt.kernels.wise_flash import compact_page_tables + + flash = _flash(generation, page) + q, k, v, tables, lengths = _pool(page=page, length=900) + rows, heads = q.shape[:2] + blocks = 1 << math.ceil(math.log2(math.ceil(tables.shape[1] * page / 32))) + support = torch.ones((rows, 1, heads, blocks), device="cuda", dtype=torch.uint8) + modes = torch.full((rows,), MODE_SPARSE, device="cuda", dtype=torch.int32) + slots = torch.arange(rows, device="cuda", dtype=torch.int32) + out_t, out_l = compact_page_tables( + tables, + lengths, + modes, + slots, + support, + 0, + page, + torch.empty_like(tables), + torch.empty_like(lengths), + ) + live = torch.arange(tables.shape[1], device="cuda")[None] * page < lengths[:, None] + assert torch.equal(out_l, lengths) and torch.equal(out_t[live], tables[live]) + assert torch.equal(flash(q, k, v, out_t, out_l), flash(q, k, v, tables, lengths)) + + +@pytest.mark.parametrize("generation,page", FLASH) +def test_flash_discovery_mass_selection_and_union_sparse_output(generation, page): + from vllm_rlt.kernels.wise_flash import compact_page_tables, discover_support + + flash = _flash(generation, page) + q, k, v, tables, lengths = _pool(page=page, length=900, seed=5) + rows, heads = q.shape[:2] + blocks = 1 << math.ceil(math.log2(math.ceil(tables.shape[1] * page / 32))) + support = torch.zeros((rows, 1, heads, blocks), device="cuda", dtype=torch.uint8) + slots = torch.arange(rows, device="cuda", dtype=torch.int32) + discover = torch.full((rows,), MODE_DISCOVER, device="cuda", dtype=torch.int32) + scratch = wise_scratch(rows, heads, q.shape[-1], blocks, "cuda") + stats = torch.zeros(2, device="cuda", dtype=torch.int64) + _, lse = flash(q, k, v, tables, lengths, return_lse=True) + discover_support(q, k, tables, lengths, lse, discover, slots, support, 0, scratch, stats) + mass = scratch["mass"][1] + valid = torch.arange(blocks, device="cuda")[None] * 32 < lengths[:, None] + normalized = mass.masked_fill(~valid[:, None], 0) + normalized = normalized / normalized.sum(-1, keepdim=True) + assert torch.equal(support[:, 0].bool(), select_support(normalized, valid, 0.95)) + for row in range(rows): + scores, values = _dense_probs(q, k, v, tables, lengths, row) + n = scores.shape[-1] + nb = math.ceil(n / 32) + sums = F.pad(scores.softmax(-1), (0, nb * 32 - n)).view(heads, nb, 32).sum(-1) + torch.testing.assert_close(normalized[row, :, :nb], sums, atol=5e-3, rtol=5e-3) + sparse = torch.full((rows,), MODE_SPARSE, device="cuda", dtype=torch.int32) + out_t, out_l = compact_page_tables( + tables, + lengths, + sparse, + slots, + support, + 0, + page, + torch.empty_like(tables), + torch.empty_like(lengths), + ) + out = flash(q, k, v, out_t, out_l) + for row in range(rows): + scores, values = _dense_probs(q, k, v, tables, lengths, row) + n = scores.shape[-1] + block_keep = support[row, 0].bool().any(0) + page_keep = ( + block_keep.repeat_interleave(32)[: math.ceil(n / page) * page].view(-1, page).any(-1) + ) + keep = page_keep.repeat_interleave(page)[:n] + weights = scores.masked_fill(~keep[None], -torch.inf).softmax(-1) + torch.testing.assert_close( + out[row].float(), torch.einsum("ht,thd->hd", weights, values), atol=2e-2, rtol=2e-2 + ) + + +@pytest.mark.parametrize("generation", [3, 4]) +def test_flash_engine_all_blocks_matches_dense_and_wise_draft_is_lossless(model, generation): + _flash(generation, 16) + prompts = [[int(t) for t in (torch.arange(300 + 50 * i) * (i + 3)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=10, ignore_eos=True) + + def engine(wise, speculative=None): + return LLM( + model, + attention_backend=f"flash_attn_{generation}", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=512, prefill_chunk_size=512 + ), + speculative_config=speculative, + wise_config=wise, + ) + + dense = [o.token_ids for o in engine(None).generate(prompts, params)] + forced = engine(WiseConfig(max_model_len=1024, max_slots=8, support_override="all")) + assert [o.token_ids for o in forced.generate(prompts, params)] == dense + real = engine(WiseConfig(max_model_len=1024, max_slots=8)) + assert all(len(o.token_ids) == 10 for o in real.generate(prompts, params)) + assert real.engine.cache_manager.wise.mean_density() < 1 + drafted = engine( + WiseConfig(max_model_len=1024, max_slots=8), SpeculativeConfig(3, draft_attention="wise") + ) + assert [o.token_ids for o in drafted.generate(prompts, params)] == dense + + +@pytest.mark.parametrize("generation", [3, 4]) +def test_triton_wise_kernel_beside_flash_prefill(model, generation): + _flash(generation, 16) + prompts = [[int(t) for t in (torch.arange(300 + 50 * i) * (i + 3)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=10, ignore_eos=True) + + def engine(override): + return LLM( + model, + attention_backend=f"flash_attn_{generation}", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=512, prefill_chunk_size=512 + ), + wise_config=WiseConfig( + max_model_len=1024, max_slots=8, support_override=override, kernel="triton" + ), + ) + + dense = engine("dense") + assert dense.engine.cache_manager.wise.flash is None + reference = [o.token_ids for o in dense.generate(prompts, params)] + assert [o.token_ids for o in engine("all").generate(prompts, params)] == reference + real = engine(None) + assert all(len(o.token_ids) == 10 for o in real.generate(prompts, params)) + assert real.engine.cache_manager.wise.mean_density() < 1 + + +@pytest.mark.parametrize("graphs", [False, True]) +def test_mixed_packed_fa4_decode_rows_fall_back_to_dense(model, graphs): + _flash(4, 16) + # Staggered lengths: later prompts' chunks share packed batches with decode rows. + prompts = [[int(t) for t in (torch.arange(200 + 150 * i) * (i + 3)) % 60 + 2] for i in range(3)] + params = SamplingParams(max_tokens=10, ignore_eos=True) + + def engine(override): + return LLM( + model, + attention_backend="flash_attn_4", + cache_config=CacheConfig(num_blocks=512, block_size=16), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=96, prefill_chunk_size=64, mode="mixed" + ), + execution_config=ExecutionConfig(atomic_prefill=False, cuda_graphs=graphs), + wise_config=WiseConfig(max_model_len=1024, max_slots=8, support_override=override), + ) + + reference = [o.token_ids for o in engine("dense").generate(prompts, params)] + forced = engine("all") + assert [o.token_ids for o in forced.generate(prompts, params)] == reference + assert forced.engine.cache_manager.wise.fallbacks > 0 + real = engine(None) + assert all(len(o.token_ids) == 10 for o in real.generate(prompts, params)) diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 648781d7..4eb4ff9a 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -180,6 +180,10 @@ class SpeculativeConfig: draft_loops: int = 2 target_loops: int = 4 mode: str = "block" + # "truncate": draft reads out after draft_loops. "wise": draft runs every + # target loop, with WISE working-set attention after draft_loops; the + # verification recomputes those loops densely from the shared draft_loops state. + draft_attention: str = "truncate" def __post_init__(self): for name in ("num_speculative_tokens", "draft_loops", "target_loops"): @@ -190,3 +194,7 @@ def __post_init__(self): raise ValueError("pipeline draft_loops must divide target_loops") if self.draft_loops >= self.target_loops: raise ValueError("draft_loops must be smaller than target_loops") + if self.draft_attention not in ("truncate", "wise"): + raise ValueError("draft_attention must be 'truncate' or 'wise'") + if self.draft_attention == "wise" and self.mode != "block": + raise ValueError("draft_attention='wise' requires block speculation") diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..7a15ba01 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -155,6 +155,9 @@ def __init__( self._prefixes = OrderedDict() self._pending_prefixes = [] self.prefix_hits = self.prefix_queries = 0 + # Optional WISE decode state (vllm_rlt.wise). The runner marks decode + # metadata; unmarked traversals, including prefill, stay dense. + self.wise = None if enable_prefix_caching and layout != "last_exited": raise ValueError("prefix caching requires last_exited KV") @@ -363,6 +366,8 @@ def free(self, request_id: str) -> None: allocation = self._allocations.pop(request_id, None) if allocation is not None: self._drop_refs(b for table in allocation.block_tables for b in reversed(table)) + if self.wise is not None: + self.wise.release(request_id) def pin_transfer(self, request_id: str, transfer_id: str): allocation = self._get_allocation(request_id) @@ -740,6 +745,10 @@ def _attend_prepared( batch.cu_seqlens_q, batch.max_seqlen_q, ) + if self.wise is not None and self.wise.is_decode(batch): + return self.wise.attend( + layer, batch, q, self.key_cache[:, layer], self.value_cache[:, layer] + ) return self.attention( q, self.key_cache[:, layer], diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..70fc42bd 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -53,6 +53,7 @@ def __init__( speculative_config=None, stage_config=None, connector: StageConnector | None = None, + wise_config=None, ): self.model = model self.stage_config = stage_config or StageConfig() @@ -195,6 +196,10 @@ def __init__( watermark_ratio=cache_config.watermark_ratio, recurrent_layers=getattr(model, "recurrent_kv_layers", None), ) + if wise_config is not None: + self._init_wise(wise_config, speculative_config, config, parameter) + elif speculative_config is not None and speculative_config.draft_attention == "wise": + raise ValueError("draft_attention='wise' requires wise_config") self.scheduler = Scheduler( scheduler_config, self.cache_manager, @@ -263,6 +268,40 @@ def __init__( self.scheduler.suspend_callback = self.preemption.suspend self.scheduler.resume_callback = self.preemption.resume + def _init_wise(self, wise_config, speculative_config, config, parameter): + """WISE working-set decode (vllm_rlt.wise) on this engine's cache.""" + from vllm_rlt.wise import WiseState + + backend = self.cache_manager.backend + if backend not in ("triton", "flash_attn_2", "flash_attn_3", "flash_attn_4"): + raise ValueError("WISE decode requires Triton or FlashAttention 2/3/4 attention") + if max(wise_config.reuse_depths) >= config.total_ut_steps: + raise ValueError("WISE reuse depths exceed the model's recurrent depth") + if speculative_config is not None: + if self.async_speculative: + raise ValueError("a WISE draft supports synchronous speculative decoding only") + if ( + speculative_config.draft_attention != "wise" + or tuple(range(speculative_config.draft_loops)) != wise_config.discovery_depths + or tuple(range(speculative_config.draft_loops, speculative_config.target_loops)) + != wise_config.reuse_depths + ): + raise ValueError( + "with speculation, WISE is the draft only: draft_attention='wise' and " + "draft/target loops must match the discovery/reuse depths" + ) + wise = WiseState( + wise_config, + config.num_hidden_layers, + config.num_attention_heads, + parameter.device, + head_dim=config.head_dim, + ) + wise.draft_only = speculative_config is not None + if backend != "triton" and wise_config.kernel == "auto": + wise.flash = self.cache_manager.attention.kernel + self.cache_manager.wise = wise + @property def connector(self): return self._connector diff --git a/vllm_rlt/entrypoints/llm.py b/vllm_rlt/entrypoints/llm.py index 55e46104..ee198649 100644 --- a/vllm_rlt/entrypoints/llm.py +++ b/vllm_rlt/entrypoints/llm.py @@ -28,6 +28,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + wise_config=None, ): tokenizer_source = resolve_model_source(model) if isinstance(model, str) else None model, revision, model_config = prepare_model(model, revision=revision) @@ -50,6 +51,7 @@ def __init__( exit_config=exit_config, execution_config=execution_config, speculative_config=speculative_config, + wise_config=wise_config, ) self._next_request_id = 0 diff --git a/vllm_rlt/kernels/flash_attention.py b/vllm_rlt/kernels/flash_attention.py index 8555464b..c7e1a717 100644 --- a/vllm_rlt/kernels/flash_attention.py +++ b/vllm_rlt/kernels/flash_attention.py @@ -95,9 +95,11 @@ def prefill( ) return out[0] if isinstance(out, tuple) else out - def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): + def __call__(self, q, key_cache, value_cache, block_tables, context_lengths, return_lse=False): + """Decode attention; with ``return_lse`` also the [rows, heads] natural-log normalizers.""" if q.shape[0] == 0: - return torch.empty_like(q) + empty = torch.empty_like(q) + return (empty, q.new_empty(q.shape[:2], dtype=torch.float32)) if return_lse else empty # The visible prefix already includes exactly the current causal position. # Do not append K/V here: the cache manager has written it at this depth. query = q.unsqueeze(1) @@ -114,6 +116,7 @@ def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): # Auto splitting depends on batch size and changes BF16 reductions. # Keep the partition fixed across scheduler batch compositions. num_splits=1, + **({"return_lse": True} if return_lse else {}), ) else: page_argument = "block_table" if self.generation == 2 else "page_table" @@ -124,7 +127,18 @@ def __call__(self, q, key_cache, value_cache, block_tables, context_lengths): cache_seqlens=context_lengths, causal=True, **{page_argument: block_tables}, + **({"return_softmax_lse": True} if return_lse else {}), ) + lse = None if isinstance(out, tuple): - out = out[0] # FA4 beta returns (output, optional LSE). + out, lse = out[0], out[1] # FA4 beta returns (output, optional LSE). + if return_lse: + # [rows, heads, 1] or [rows, 1, heads]; one query position per row. + rows, heads = q.shape[:2] + lse = ( + lse.reshape(rows, heads) + if lse.shape[1] == heads + else lse.transpose(1, 2).reshape(rows, heads) + ) + return out.squeeze(1), lse.float() return out.squeeze(1) diff --git a/vllm_rlt/kernels/wise_attention.py b/vllm_rlt/kernels/wise_attention.py new file mode 100644 index 00000000..fae78ed1 --- /dev/null +++ b/vllm_rlt/kernels/wise_attention.py @@ -0,0 +1,445 @@ +"""Split-KV paged decode attention with WISE working-set support (arXiv 2609.27373). + +Two launches per layer call: + +* ``_wise_split_kernel`` over [row, query head, split]: each split attends a + contiguous share of the row's 32-token tiles (one WISE block per tile) and + writes a partial online-softmax state. Per row mode: + + - ``MODE_DENSE``: every tile; + - ``MODE_DISCOVER`` / ``MODE_DISCOVER_UNION``: every tile, also recording each + tile's (max, sum exp); + - ``MODE_SPARSE``: only tiles whose support bit is set, compacted first so the + attention loop is branch-free. Splits divide the selected tiles, so a sparse + row's work, not just its bytes, shrinks with density. + +* ``_wise_finish_kernel`` over [row, query head]: merges the partial states and, + for discovery rows, rebuilds per-block mass and selects the smallest stable + descending-mass prefix reaching ``eta`` (the authors' direct_block_support rule + for one query row). The first discovery depth overwrites the support bitmask, + later ones OR into it. + +Splitting keeps one long dense row from setting the latency of a launch whose +other rows are sparse. With every support bit set a sparse row uses the same +tiles, splits and arithmetic as MODE_DENSE, so its output is bitwise identical. +""" + +import math + +import torch +import triton +import triton.language as tl + +MODE_DENSE, MODE_DISCOVER, MODE_SPARSE, MODE_DISCOVER_UNION = 0, 1, 2, 3 +WISE_BLOCK = 32 +MAX_SPLITS = 32 + + +def num_splits(rows, heads, device, max_tiles): + """Splits per [row, head]: about eight 32-token tiles per program for the longest + row, and at least four programs per SM. A dense program scanning thousands of + tokens is latency-bound and sets the time of a launch whose other rows are sparse.""" + sms = torch.cuda.get_device_properties(device).multi_processor_count + wanted = max(math.ceil(max_tiles / 8), math.ceil(4 * sms / max(1, rows * heads)), 1) + return min(MAX_SPLITS, triton.next_power_of_2(wanted)) + + +@triton.jit +def _wise_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + MODES, + SLOTS, + SUPPORT, + MASS_M, + MASS_S, + SELECTED, + PART_M, + PART_L, + PART_ACC, + layer, + q_batch_stride, + q_head_stride, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + table_stride, + support_slot_stride, + support_layer_stride, + support_head_stride, + mass_row_stride, + mass_head_stride, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + HEADS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + MAX_BLOCKS: tl.constexpr, + SPLITS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + kv_head = head // HEAD_GROUPS + length = tl.load(LENGTHS + row) + mode = tl.load(MODES + row) + mass = row.to(tl.int64) * mass_row_stride + head * mass_head_stride + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_batch_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + count = (length + BLOCK_T - 1) // BLOCK_T + total = count + if mode == 2: + slot = tl.load(SLOTS + row).to(tl.int64) + support = ( + SUPPORT + + slot * support_slot_stride + + layer * support_layer_stride + + head * support_head_stride + ) + index = tl.arange(0, MAX_BLOCKS) + bits = tl.load(support + index, mask=index < count, other=0) != 0 + target = tl.cumsum(bits.to(tl.int32), axis=0) - 1 + # Every split writes the same list for its [row, head], then reads it back. + tl.store(SELECTED + mass + target, index, mask=bits) + total = tl.sum(bits.to(tl.int32), axis=0) + tl.debug_barrier() + share = (total + SPLITS - 1) // SPLITS + first = split * share + last = tl.minimum(first + share, total) + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + sparse = mode == 2 + record = (mode == 1) | (mode == 3) + for item in range(first, last): + # Masked loads/stores instead of branches keep the loop pipelinable. + tile = tl.where(sparse, tl.load(SELECTED + mass + item, mask=sparse, other=0), item) + positions = tile * BLOCK_T + tl.arange(0, BLOCK_T) + valid_tokens = positions < length + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, mask=valid_tokens, other=0 + ).to(tl.int64) + offsets = positions % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) * SCALE + scores = tl.where(valid_tokens, scores, -float("inf")) + tile_maximum = tl.max(scores, axis=0) + next_maximum = tl.maximum(maximum, tile_maximum) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + tl.store(MASS_M + mass + tile, tile_maximum, mask=record) + tl.store(MASS_S + mass + tile, tl.sum(tl.exp(scores - tile_maximum), axis=0), mask=record) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) + normalizer = normalizer * correction + tl.sum(weights, axis=0) + maximum = next_maximum + part = (row.to(tl.int64) * HEADS + head) * SPLITS + split + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_ACC + part * BLOCK_D + dims, accumulator) + + +@triton.jit +def _wise_finish_kernel( + LENGTHS, + MODES, + SLOTS, + SUPPORT, + MASS_M, + MASS_S, + STATS, + PART_M, + PART_L, + PART_ACC, + OUT, + layer, + out_batch_stride, + out_head_stride, + support_slot_stride, + support_layer_stride, + support_head_stride, + mass_row_stride, + mass_head_stride, + HEAD_DIM: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + MAX_BLOCKS: tl.constexpr, + SPLITS: tl.constexpr, + ETA: tl.constexpr, + COMBINE: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + mode = tl.load(MODES + row) + if COMBINE: + dims = tl.arange(0, BLOCK_D) + splits = tl.arange(0, SPLITS) + part = (row.to(tl.int64) * HEADS + head) * SPLITS + maxima = tl.load(PART_M + part + splits) + normalizers = tl.load(PART_L + part + splits) + accumulators = tl.load(PART_ACC + (part + splits)[:, None] * BLOCK_D + dims[None, :]) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima == -float("inf"), 0.0, tl.exp(maxima - maximum)) + normalizer = tl.sum(normalizers * scale, axis=0) + accumulator = tl.sum(accumulators * scale[:, None], axis=0) + tl.store( + OUT + row * out_batch_stride + head * out_head_stride + dims, + accumulator / tl.maximum(normalizer, 1.0e-30), + mask=dims < HEAD_DIM, + ) + if (mode == 1) | (mode == 3): + length = tl.load(LENGTHS + row) + slot = tl.load(SLOTS + row).to(tl.int64) + support = ( + SUPPORT + + slot * support_slot_stride + + layer * support_layer_stride + + head * support_head_stride + ) + mass = row.to(tl.int64) * mass_row_stride + head * mass_head_stride + index = tl.arange(0, MAX_BLOCKS) + count = (length + BLOCK_T - 1) // BLOCK_T + valid = index < count + tile_m = tl.load(MASS_M + mass + index, mask=valid, other=-float("inf")) + tile_s = tl.load(MASS_S + mass + index, mask=valid, other=0.0) + weight = tile_s * tl.exp(tile_m - tl.max(tile_m, axis=0)) + weight = tl.where(valid, weight, 0.0) + probability = weight / tl.sum(weight, axis=0) + # Stable descending sort: nonnegative float bits are monotone as int32; + # the low half breaks ties toward the lower block index. Blocks past + # the row sort last (negative keys) but keep their index in the low + # half, so the store below also clears their stale bits. + bits = probability.to(tl.int32, bitcast=True).to(tl.int64) + low = (MAX_BLOCKS - 1 - index).to(tl.int64) + key = tl.where(valid, (bits << 32) | low, low - (1 << 62)) + ranked = tl.sort(key, descending=True) + order = MAX_BLOCKS - 1 - (ranked & 0xFFFFFFFF) + ranked_mass = tl.where( + ranked >= 0, ((ranked >> 32).to(tl.int32)).to(tl.float32, bitcast=True), 0.0 + ) + total = tl.sum(ranked_mass, axis=0) + below = tl.sum((tl.cumsum(ranked_mass, axis=0) < ETA * total).to(tl.int32), axis=0) + keep_count = tl.maximum(tl.minimum(below + 1, count), 1) + chosen = (index < keep_count).to(tl.uint8) + # STATS counts the support the reuse depths attend: the first discovery + # adds its selection and causal blocks, a union adds only the blocks it sets. + if mode == 3: + previous = tl.load(support + order) + chosen = chosen | previous + tl.atomic_add(STATS, tl.sum((chosen - previous).to(tl.int64), axis=0)) + else: + tl.atomic_add(STATS, keep_count.to(tl.int64)) + tl.atomic_add(STATS + 1, count.to(tl.int64)) + tl.store(support + order, chosen) + + +def wise_scratch(rows, heads, head_dim, max_blocks, device): + """Per-traversal scratch: tile mass (also the compacted tile lists) and split partials.""" + width = triton.next_power_of_2(head_dim) + return { + "mass": torch.empty((2, rows, heads, max_blocks), device=device, dtype=torch.float32), + "part_m": torch.empty((rows, heads, MAX_SPLITS), device=device, dtype=torch.float32), + "part_l": torch.empty((rows, heads, MAX_SPLITS), device=device, dtype=torch.float32), + "part_acc": torch.empty( + (rows, heads, MAX_SPLITS, width), device=device, dtype=torch.float32 + ), + } + + +def select_from_mass( + context_lengths, modes, slots, support, layer, mass_m, mass_s, stats, eta, heads +): + """Discovery selection only (FlashAttention path): masses are already recorded.""" + rows = context_lengths.shape[0] + if rows == 0: + return + _wise_finish_kernel[(rows, heads)]( + context_lengths, + modes, + slots, + support, + mass_m, + mass_s, + stats, + mass_m, + mass_m, + mass_m, + mass_m, + layer, + 0, + 0, + support.stride(0), + support.stride(1), + support.stride(2), + mass_m.stride(0), + mass_m.stride(1), + HEAD_DIM=1, + HEADS=heads, + BLOCK_D=1, + BLOCK_T=WISE_BLOCK, + MAX_BLOCKS=support.shape[-1], + SPLITS=1, + ETA=eta, + COMBINE=False, + num_warps=4, + ) + + +def wise_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + modes, + slots, + support, + layer, + scratch, + stats, + eta=0.95, + splits=None, + num_warps=1, + num_stages=None, +): + """Split-KV WISE decode over [row, query head]. + + One warp per split program measured fastest on GH200 (4K x 20 dense: 207 us at + 1 warp, 312 us at 4, FlashAttention-3 211 us). + + ``support`` is [slots, layers, query heads, MAX_BLOCKS] uint8, MAX_BLOCKS a power + of two covering every row. ``scratch`` comes from ``wise_scratch``; ``stats`` is + int64 [2] accumulating (selected, causal) block counts over discovery rows. + """ + output = torch.empty_like(q) + rows, heads = q.shape[:2] + if rows == 0: + return output + splits = splits or num_splits(rows, heads, q.device, support.shape[-1]) + mass = scratch["mass"] + block_d = triton.next_power_of_2(q.shape[-1]) + _wise_split_kernel[(rows, heads, splits)]( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + modes, + slots, + support, + mass[0], + mass[1], + mass[1].view(torch.int32), + scratch["part_m"], + scratch["part_l"], + scratch["part_acc"], + layer, + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + block_tables.stride(0), + support.stride(0), + support.stride(1), + support.stride(2), + mass.stride(1), + mass.stride(2), + HEAD_DIM=q.shape[-1], + HEAD_GROUPS=heads // key_cache.shape[2], + HEADS=heads, + PAGE_SIZE=key_cache.shape[1], + SCALE=q.shape[-1] ** -0.5, + BLOCK_D=block_d, + BLOCK_T=WISE_BLOCK, + MAX_BLOCKS=support.shape[-1], + SPLITS=splits, + num_warps=num_warps, + **({"num_stages": num_stages} if num_stages else {}), + ) + _wise_finish_kernel[(rows, heads)]( + context_lengths, + modes, + slots, + support, + mass[0], + mass[1], + stats, + scratch["part_m"], + scratch["part_l"], + scratch["part_acc"], + output, + layer, + output.stride(0), + output.stride(1), + support.stride(0), + support.stride(1), + support.stride(2), + mass.stride(1), + mass.stride(2), + HEAD_DIM=q.shape[-1], + HEADS=heads, + BLOCK_D=block_d, + BLOCK_T=WISE_BLOCK, + MAX_BLOCKS=support.shape[-1], + SPLITS=splits, + ETA=eta, + COMBINE=True, + num_warps=4, + ) + return output + + +def block_mass(mass_m, mass_s, lengths): + """Reference per-block attention probability mass [rows, heads, blocks] from tile records.""" + blocks = torch.arange(mass_m.shape[-1], device=mass_m.device) + valid = blocks[None, :] * WISE_BLOCK < lengths[:, None] + m = mass_m.masked_fill(~valid[:, None, :], -float("inf")) + s = mass_s.masked_fill(~valid[:, None, :], 0.0) + weights = s * torch.exp(m - m.amax(-1, keepdim=True)) + return weights / weights.sum(-1, keepdim=True), valid + + +def select_support(mass, valid, eta): + """Reference rule: smallest stable descending-mass prefix reaching ``eta`` of the causal mass. + + Same rule as the pinned author ``direct_block_support`` for one query row. + """ + ranked, order = torch.sort( + mass.masked_fill(~valid[:, None, :], -float("inf")), descending=True, stable=True + ) + total = mass.sum(-1, keepdim=True) + keep = ranked.clamp_min(0).cumsum(-1).lt(eta * total).sum(-1) + 1 + keep = torch.minimum(keep, valid.sum(-1, keepdim=True).expand_as(keep)).clamp_min(1) + chosen = torch.arange(mass.shape[-1], device=mass.device) < keep.unsqueeze(-1) + return torch.zeros_like(chosen).scatter_(-1, order, chosen) & valid[:, None, :] diff --git a/vllm_rlt/kernels/wise_flash.py b/vllm_rlt/kernels/wise_flash.py new file mode 100644 index 00000000..1bbf992c --- /dev/null +++ b/vllm_rlt/kernels/wise_flash.py @@ -0,0 +1,223 @@ +"""WISE working-set decode on top of FlashAttention 2/3/4 paged kernels. + +FlashAttention takes one page table per row for all heads, so sparse rows +attend the union over heads of their WISE block support, at page granularity +(a page is kept if any of its 32-token blocks is selected by any head). +``compact_page_tables`` rewrites each sparse row's page table to only those +pages, keeping order so that the partially filled newest page stays last; +dense and discovery rows keep their table. With every block selected the +compacted table and length equal the original ones, so FlashAttention's output +is unchanged bit for bit. + +Discovery rows run dense FlashAttention with its log-sum-exp; ``discover_support`` +then rereads K once (split across programs like the attention), rebuilds each +32-token block's probability mass from the LSE, and applies the same stable +smallest-prefix selection as the Triton kernel. +""" + +import triton +import triton.language as tl + +from vllm_rlt.kernels.wise_attention import WISE_BLOCK, num_splits, select_from_mass + + +@triton.jit +def _compact_kernel( + TABLES, + LENGTHS, + MODES, + SLOTS, + SUPPORT, + OUT_TABLES, + OUT_LENGTHS, + layer, + table_stride, + out_table_stride, + support_slot_stride, + support_layer_stride, + support_head_stride, + HEADS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK: tl.constexpr, + MAX_PAGES: tl.constexpr, + BLOCKS_PER_PAGE: tl.constexpr, +): + row = tl.program_id(0) + length = tl.load(LENGTHS + row) + mode = tl.load(MODES + row) + pages = tl.arange(0, MAX_PAGES) + count = (length + PAGE_SIZE - 1) // PAGE_SIZE + live = pages < count + table = tl.load(TABLES + row * table_stride + pages, mask=live, other=0) + if mode == 2: + slot = tl.load(SLOTS + row).to(tl.int64) + base = SUPPORT + slot * support_slot_stride + layer * support_layer_stride + keep = tl.zeros((MAX_PAGES,), dtype=tl.int32) + for head in range(HEADS): + for part in range(BLOCKS_PER_PAGE): + blocks = (pages * PAGE_SIZE) // BLOCK + part + bits = tl.load(base + head * support_head_stride + blocks, mask=live, other=0) + keep = keep | bits.to(tl.int32) + keep = keep & live.to(tl.int32) + target = tl.cumsum(keep, axis=0) - 1 + tl.store(OUT_TABLES + row * out_table_stride + target, table, mask=keep != 0) + kept = tl.sum(keep, axis=0) + newest = tl.sum(tl.where(pages == count - 1, keep, 0), axis=0) + tail = length - (count - 1) * PAGE_SIZE + tl.store(OUT_LENGTHS + row, (kept - newest) * PAGE_SIZE + newest * tail) + else: + tl.store(OUT_TABLES + row * out_table_stride + pages, table, mask=live) + tl.store(OUT_LENGTHS + row, length) + + +def compact_page_tables( + block_tables, context_lengths, modes, slots, support, layer, page_size, out_tables, out_lengths +): + rows = block_tables.shape[0] + if rows == 0: + return out_tables[:0], out_lengths[:0] + max_pages = triton.next_power_of_2(block_tables.shape[1]) + _compact_kernel[(rows,)]( + block_tables, + context_lengths, + modes, + slots, + support, + out_tables, + out_lengths, + layer, + block_tables.stride(0), + out_tables.stride(0), + support.stride(0), + support.stride(1), + support.stride(2), + HEADS=support.shape[2], + PAGE_SIZE=page_size, + BLOCK=WISE_BLOCK, + MAX_PAGES=max_pages, + BLOCKS_PER_PAGE=max(1, page_size // WISE_BLOCK), + num_warps=4, + ) + return out_tables[:rows], out_lengths[:rows] + + +@triton.jit +def _flash_mass_kernel( + Q, + K, + TABLES, + LENGTHS, + LSE, + MODES, + MASS_M, + MASS_S, + q_batch_stride, + q_head_stride, + k_block_stride, + k_token_stride, + k_head_stride, + table_stride, + lse_batch_stride, + lse_head_stride, + mass_row_stride, + mass_head_stride, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, +): + """Discovery rows: each split records exp(score - LSE) summed per tile (max 0).""" + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + mode = tl.load(MODES + row) + if (mode == 1) | (mode == 3): + kv_head = head // HEAD_GROUPS + length = tl.load(LENGTHS + row) + lse = tl.load(LSE + row * lse_batch_stride + head * lse_head_stride) + mass = row.to(tl.int64) * mass_row_stride + head * mass_head_stride + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_batch_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + count = (length + BLOCK_T - 1) // BLOCK_T + share = (count + SPLITS - 1) // SPLITS + for tile in range(split * share, tl.minimum(split * share + share, count)): + positions = tile * BLOCK_T + tl.arange(0, BLOCK_T) + valid_tokens = positions < length + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, mask=valid_tokens, other=0 + ).to(tl.int64) + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + (positions % PAGE_SIZE)[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) * SCALE + weights = tl.where(valid_tokens, tl.exp(scores - lse), 0.0) + tl.store(MASS_M + mass + tile, 0.0) + tl.store(MASS_S + mass + tile, tl.sum(weights, axis=0)) + + +def discover_support( + q, + key_cache, + block_tables, + context_lengths, + lse, + modes, + slots, + support, + layer, + scratch, + stats, + eta=0.95, +): + """Selection for discovery rows from FlashAttention's LSE; other rows are untouched. + + ``lse`` is [rows, heads] natural-log normalizers of the dense attention just run. + ``scratch`` comes from ``wise_scratch``; its mass planes are left holding the masses. + """ + rows, heads = q.shape[:2] + if rows == 0: + return + mass = scratch["mass"] + splits = num_splits(rows, heads, q.device, support.shape[-1]) + _flash_mass_kernel[(rows, heads, splits)]( + q, + key_cache, + block_tables, + context_lengths, + lse, + modes, + mass[0], + mass[1], + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + block_tables.stride(0), + lse.stride(0), + lse.stride(1), + mass.stride(1), + mass.stride(2), + HEAD_DIM=q.shape[-1], + HEAD_GROUPS=heads // key_cache.shape[2], + PAGE_SIZE=key_cache.shape[1], + SCALE=q.shape[-1] ** -0.5, + BLOCK_D=triton.next_power_of_2(q.shape[-1]), + BLOCK_T=WISE_BLOCK, + SPLITS=splits, + num_warps=1, + ) + select_from_mass( + context_lengths, modes, slots, support, layer, mass[0], mass[1], stats, eta, heads + ) diff --git a/vllm_rlt/wise.py b/vllm_rlt/wise.py new file mode 100644 index 00000000..fa4466f0 --- /dev/null +++ b/vllm_rlt/wise.py @@ -0,0 +1,332 @@ +"""WISE working-set decode for recurrent-depth models (Wan and Chen, arXiv 2609.27373). + +Variant ``wise-cached-decode``: during decode, the query row of each token +discovers block support with unrestricted attention at the discovery depths, +the union of those supports is reused at the later depths, and Q/K/V plus the +within-support weights stay dynamic. Differences from the paper's no-cache +reference: the cached K/V of earlier tokens were computed under their own +supports, and the query block is the single decode row. Prefill stays dense. + +Support is request-local: it is tagged with (request, position), rebuilt for +every token, and a reuse depth never uses support from another position or an +incomplete discovery window. When a token's discovery is not complete, its +remaining depths attend densely instead: discovery depths that ran in a packed +batch (prefill attention records no support), or a KV offload between discovery +and reuse (``free`` releases the support). ``fallbacks`` counts those rows. +""" + +import math +import weakref +from dataclasses import dataclass + +import torch + +from vllm_rlt.kernels.wise_attention import ( + MODE_DENSE, + MODE_DISCOVER, + MODE_DISCOVER_UNION, + MODE_SPARSE, + WISE_BLOCK, + wise_paged_attention, + wise_scratch, +) +from vllm_rlt.kernels.wise_flash import compact_page_tables, discover_support + + +@dataclass(frozen=True) +class WiseConfig: + profile: str = "ouro4" + discovery_depths: tuple[int, ...] = (0, 1) + reuse_depths: tuple[int, ...] = (2, 3) + mass_threshold: float = 0.95 + max_model_len: int = 8192 + max_slots: int = 256 + # Benchmark controls only: "all" forces every block on at reuse depths while + # discovery still runs; "dense" runs every row densely through the same split-KV + # kernel; a float in (0, 1] skips discovery and keeps a fixed fraction of blocks + # (always including the newest one). + support_override: str | float | None = None + # "auto": WISE decode on the cache's attention backend. "triton": the Triton + # split-KV WISE kernel even when FlashAttention serves prefill and dense decode. + kernel: str = "auto" + + def __post_init__(self): + if set(self.discovery_depths) & set(self.reuse_depths) or not self.discovery_depths: + raise ValueError("discovery and reuse depths must be disjoint and nonempty") + if max(self.discovery_depths) >= min(self.reuse_depths, default=math.inf): + raise ValueError("every reuse depth must follow the discovery window") + if not 0 < self.mass_threshold <= 1: + raise ValueError("mass_threshold must be in (0, 1]") + if self.kernel not in ("auto", "triton"): + raise ValueError("kernel must be 'auto' or 'triton'") + if self.support_override not in (None, "all", "dense") and not ( + isinstance(self.support_override, float) and 0 < self.support_override <= 1 + ): + raise ValueError( + "support_override must be None, 'all', 'dense', or a fraction in (0, 1]" + ) + + +class WiseState: + """Per-request support bitmasks [slot, layer, query head, block] on the device.""" + + def __init__( + self, config: WiseConfig, num_layers: int, num_heads: int, device, head_dim: int = 128 + ): + self.config = config + self.head_dim = head_dim + # Power of two: the fused selection sorts one row of blocks in registers. + self.blocks = 1 << max( + 0, math.ceil(math.log2(math.ceil(config.max_model_len / WISE_BLOCK))) + ) + self.support = torch.zeros( + (config.max_slots, num_layers, num_heads, self.blocks), device=device, dtype=torch.uint8 + ) + self.slots: dict[str, int] = {} + self.free_slots = list(reversed(range(config.max_slots))) + # request -> (position, discovery depths completed for that position); a + # position whose discovery window was broken has depths None (dense to the end). + self.tags: dict[str, tuple[int, frozenset | None]] = {} + # Decode rows that attended densely because their discovery was incomplete. + self.fallbacks = 0 + self._batch = None + self._plan = None + # Set when WISE only drafts for speculative decoding: target decode stays dense. + self.draft_only = False + # KV metadata -> per-row decode flags, set by the runner for decode traversals. + # Rows without a flag (prefill chunks, padding) attend densely. + self._decode = weakref.WeakKeyDictionary() + # A FlashPagedAttention kernel when the cache uses FlashAttention 2/3/4: sparse + # rows then attend the per-layer head union of their support (see wise_flash). + self.flash = None + # (selected, causal) block counts over all selections, for diagnostics. + self.stats = torch.zeros(2, device=device, dtype=torch.int64) + self.eta = 2.0 if config.support_override == "all" else config.mass_threshold + + def release(self, request_id: str) -> None: + slot = self.slots.pop(request_id, None) + if slot is not None: + self.free_slots.append(slot) + self.tags.pop(request_id, None) + + def _slot(self, request_id): + if request_id not in self.slots: + if not self.free_slots: + raise RuntimeError("WISE support slots exhausted; raise WiseConfig.max_slots") + self.slots[request_id] = self.free_slots.pop() + return self.slots[request_id] + + def mark(self, metadata, decode_rows): + """Declare which host rows of ``metadata`` are decode tokens (others stay dense).""" + if any(decode_rows): + self._decode[metadata] = tuple(decode_rows) + + def inherit(self, metadata, source): + """Carry decode flags to a view of ``source`` (e.g. trimmed graph metadata).""" + if source in self._decode: + self._decode[metadata] = self._decode[source] + + def is_decode(self, metadata): + return metadata in self._decode + + def plan(self, batch): + """Host bookkeeping once per traversal; reused by every layer of the batch.""" + if self._batch is batch: + return self._plan + names = {id(allocation): request_id for request_id, allocation in batch.allocations} + decode = self._decode.get(batch, ()) + modes, slots = [], [] + override = self.config.support_override + for row, (allocation, depth, position) in enumerate(batch.rows): + if row >= len(decode) or not decode[row]: + modes.append(MODE_DENSE) + slots.append(0) + continue + request_id = names[id(allocation)] + if position >= self.config.max_model_len: + raise ValueError("position exceeds WiseConfig.max_model_len") + slots.append(self._slot(request_id)) + if override == "dense": + modes.append(MODE_DENSE) + elif depth in self.config.discovery_depths and isinstance(override, float): + modes.append(MODE_DENSE) + elif depth in self.config.discovery_depths: + done = self.tags.get(request_id) + previous = done[1] if done is not None and done[0] == position else frozenset() + if depth == self.config.discovery_depths[0]: + # A new discovery window; a rejected speculative draft may revisit a position. + previous = frozenset() + expected = frozenset(d for d in self.config.discovery_depths if d < depth) + if previous != expected: + # An earlier discovery depth of this token did not record support. + self.tags[request_id] = (position, None) + self.fallbacks += 1 + modes.append(MODE_DENSE) + continue + self.tags[request_id] = (position, previous | {depth}) + first = depth == self.config.discovery_depths[0] + modes.append(MODE_DISCOVER if first else MODE_DISCOVER_UNION) + elif depth in self.config.reuse_depths: + done = self.tags.get(request_id) + if not isinstance(override, float) and ( + done is None or done != (position, frozenset(self.config.discovery_depths)) + ): + self.fallbacks += 1 + modes.append(MODE_DENSE) + continue + modes.append(MODE_SPARSE) + else: + modes.append(MODE_DENSE) + # Padding rows (no host row) attend densely and never touch support. + padding = len(batch.context_lengths) - len(modes) + modes += [MODE_DENSE] * padding + slots += [0] * padding + device = batch.context_lengths.device + modes_t = torch.tensor(modes, dtype=torch.int32, pin_memory=device.type == "cuda").to( + device, non_blocking=True + ) + slots_t = torch.tensor(slots, dtype=torch.int32, pin_memory=device.type == "cuda").to( + device, non_blocking=True + ) + plan = { + "modes": modes_t, + "slots": slots_t, + "scratch": self._scratch(len(modes), batch.block_tables, batch.context_lengths), + } + if isinstance(override, float): + sparse = [i for i, m in enumerate(modes) if m == MODE_SPARSE] + if sparse: + self._fixed_support(batch, sparse, slots, override) + self._batch, self._plan = batch, plan + return plan + + def _fixed_support(self, batch, rows, slots, fraction): + # Deterministic fixed-density control: every k-th block plus the newest. + lengths = batch.context_lengths[rows].long() + blocks = torch.arange(self.blocks, device=lengths.device) + valid = blocks[None, :] * WISE_BLOCK < lengths[:, None] + newest = (lengths[:, None] - 1) // WISE_BLOCK == blocks[None, :] + stride = max(1, round(1 / fraction)) + chosen = ((blocks[None, :] % stride == 0) | newest) & valid + index = torch.tensor([slots[i] for i in rows], device=lengths.device) + scale = self.support.shape[1] * self.support.shape[2] + self.stats[0] += chosen.sum() * scale + self.stats[1] += valid.sum() * scale + self.support[index] = ( + chosen[:, None, None, :].expand(-1, *self.support.shape[1:3], -1).to(torch.uint8) + ) + + def _scratch(self, rows, tables, lengths): + scratch = wise_scratch( + rows, self.support.shape[2], self.head_dim, self.blocks, self.support.device + ) + if self.flash is not None: + scratch["tables"], scratch["lengths"] = ( + torch.empty_like(tables), + torch.empty_like(lengths), + ) + return scratch + + def _run(self, layer, q, key_cache, value_cache, tables, lengths, modes, slots, scratch): + if self.flash is None: + return wise_paged_attention( + q, + key_cache, + value_cache, + tables, + lengths, + modes, + slots, + self.support, + layer, + scratch, + self.stats, + self.eta, + ) + sparse_tables, sparse_lengths = compact_page_tables( + tables, + lengths, + modes, + slots, + self.support, + layer, + key_cache.shape[1], + scratch["tables"], + scratch["lengths"], + ) + output, lse = self.flash( + q, key_cache, value_cache, sparse_tables, sparse_lengths, return_lse=True + ) + discover_support( + q, + key_cache, + tables, + lengths, + lse, + modes, + slots, + self.support, + layer, + scratch, + self.stats, + self.eta, + ) + return output + + def attend(self, layer, batch, q, key_cache, value_cache): + plan = self.plan(batch) + return self._run( + layer, + q, + key_cache, + value_cache, + batch.block_tables, + batch.context_lengths, + plan["modes"], + plan["slots"], + plan["scratch"], + ) + + def static_inputs(self, rows, tables, lengths): + """Graph-owned per-row modes, slots and scratch for one captured batch size.""" + device = self.support.device + return { + "modes": torch.zeros(rows, device=device, dtype=torch.int32), + "slots": torch.zeros(rows, device=device, dtype=torch.int32), + "scratch": self._scratch(rows, tables, lengths), + } + + def copy_static(self, inputs, batch): + """Host checks for this traversal, then device copies into the graph's inputs.""" + plan = self.plan(batch) + rows = len(plan["modes"]) + inputs["modes"].zero_() + inputs["modes"][:rows].copy_(plan["modes"]) + inputs["slots"][:rows].copy_(plan["slots"]) + + def attend_static(self, layer, metadata, q, key_cache, value_cache): + inputs = metadata.wise + return self._run( + layer, + q, + key_cache, + value_cache, + metadata.block_tables, + metadata.context_lengths, + inputs["modes"], + inputs["slots"], + inputs["scratch"], + ) + + def mean_density(self): + """Selected / causal blocks of the support the reuse depths attend: the union + over each token's discovery depths (or the fixed support), averaged over every + (row, layer, head) selection so far.""" + selected, causal = self.stats.tolist() + return selected / max(causal, 1) + + def density(self, request_id, layer, length): + """Diagnostics: selected fraction of causal blocks for one request and layer.""" + blocks = math.ceil(length / WISE_BLOCK) + row = self.support[self.slots[request_id], layer, :, :blocks] + return float(row.float().mean()) diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index 10149516..d9a950e7 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -121,6 +121,10 @@ def _write_prepared(self, layer, batch, k, v): def _attend_prepared(self, layer, batch, q): cache = self.cache + if getattr(batch, "wise", None) is not None: + return cache.wise.attend_static( + layer, batch, q, cache.key_cache[:, layer], cache.value_cache[:, layer] + ) if batch.cu_seqlens_q is not None: return cache.attention.prefill( q, @@ -167,10 +171,20 @@ def _graph_entry( if bucketed or getattr(batch, "write_mask", None) is not None else None ), + wise=None, ), ) +def _attach_wise(entry, cache): + """Give a WISE decode graph its own static modes, slots and scratch.""" + metadata = entry.metadata + metadata.wise = cache.wise.static_inputs( + len(metadata.context_lengths), metadata.block_tables, metadata.context_lengths + ) + return entry + + def _copy_graph_inputs(entry, hidden, batch, *, count=None): count = _query_count(batch) if count is None else count metadata = entry.metadata @@ -223,18 +237,23 @@ def __init__(self, model, cache, execution, compute_gate): def select(self, hidden, batch): count = _query_count(batch) packed = batch.cu_seqlens_q is not None + wise = self.cache.wise is not None and not packed and self.cache.wise.is_decode(batch) if len(batch.position_ids) != count: - batch = replace( + trimmed = replace( batch, position_ids=batch.position_ids[:count], context_lengths=batch.context_lengths[:count], block_tables=batch.block_tables[:count], ) + if wise: + self.cache.wise.inherit(trimmed, batch) + batch = trimmed + # WISE and dense decode capture different attention; never replay one as the other. key = ( count, len(batch.context_lengths), batch.max_seqlen_q if packed else 1, - packed, + "wise" if wise else packed, getattr(batch, "write_mask", None) is not None, ) frontier = {} @@ -260,6 +279,8 @@ def select(self, hidden, batch): entry = _graph_entry( hidden, batch, self.cache, self.model.config.max_position_embeddings ) + if wise: + _attach_wise(entry, self.cache) entry.graph = None self.entries[key] = entry return PreparedGraph(batch, entry) @@ -286,6 +307,8 @@ def prepare_inputs(self, call, hidden): if self.last_event is not None: stream.wait_event(self.last_event) _copy_graph_inputs(entry, hidden, batch) + if entry.metadata.wise is not None: + cache.wise.copy_static(entry.metadata.wise, batch) if entry.graph is None: proxy = _DeviceCache(cache) entry.graph, entry.output = _capture( diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..2669e449 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1027,6 +1027,15 @@ def _prepare_forward(self, prepared): kv = prepared.kv if kv is not None: metadata = (kv,) + metadata + wise = self.cache_manager.wise + if wise is not None and batch.emit_gate and not wise.draft_only: + # Decode rows use WISE; prefill-chunk rows mixed into the batch stay dense. + decode_rows = [ + flag for span in batch.spans for flag in [span.emit_gate] * span.token_count + ] + for kv_metadata in metadata: + if hasattr(kv_metadata, "rows"): + wise.mark(kv_metadata, decode_rows) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden[: batch.num_rows] diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 85f0e235..34a93714 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -197,12 +197,19 @@ def metadata(ids, positions, depths, *, packed=False): layouts, lambda values, dtype: self.input_buffers[dtype].append(values) ) + wise = self.config.draft_attention == "wise" steps = [] for offset in range(max(span.token_count for span in requests)): active = [i for i, span in enumerate(requests) if offset < span.token_count] ids = [requests[i].request.request_id for i in active] positions = [requests[i].token_start + offset for i in active] - steps.append((active, metadata(ids, positions, range(self.config.draft_loops)))) + # A WISE draft also runs the remaining loops (sparse) for its readout. + reuse = ( + metadata(ids, positions, range(self.config.draft_loops, self.config.target_loops)) + if wise + else () + ) + steps.append((active, metadata(ids, positions, range(self.config.draft_loops)), reuse)) ids, positions = [], [] for span in requests: ids.extend([span.request.request_id] * span.token_count) @@ -215,9 +222,18 @@ def metadata(ids, positions, depths, *, packed=False): tokens.shape, tokens.stride(), tokens.storage_offset() ) steps = [ - (active, tuple(self.metadata.bind(kv, uploaded) for kv in layouts)) - for active, layouts in steps + ( + active, + tuple(self.metadata.bind(kv, uploaded) for kv in layouts), + tuple(self.metadata.bind(kv, uploaded) for kv in reuse), + ) + for active, layouts, reuse in steps ] + if wise: + # Every draft row is a WISE decode row; verification stays dense. + for active, layouts, reuse in steps: + for kv in layouts + reuse: + self.cache.wise.mark(kv, [True] * len(active)) verification = tuple(self.metadata.bind(kv, uploaded) for kv in verification) return tokens, steps, verification @@ -423,7 +439,7 @@ def _submit_round(self, batch, inputs, steps, verification): # prelude input, and all of them are read back once after verification. drafted = [] # Each step consumes the previous step's GPU-resident sampled IDs. - for offset, (active, metadata) in enumerate(steps): + for offset, (active, metadata, reuse) in enumerate(steps): hidden = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=inputs)).hidden hidden = self._recurrent(hidden, metadata) drafting = [] @@ -433,6 +449,10 @@ def _submit_round(self, batch, inputs, steps, verification): drafting.append((row, i)) if not drafting: break + if reuse: + # Shared dense discovery loops above; WISE-sparse remaining loops + # only for the draft readout. Verification overwrites their KV. + hidden = self._recurrent(hidden, reuse) logits = self._coda(torch.stack([hidden[row] for row, _ in drafting])) inputs = logits.argmax(-1) for row, (_, i) in enumerate(drafting):