diff --git a/README.md b/README.md index 00ed4e17..71dc3be7 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,10 @@ currently **same-host CUDA** with supported Triton/FA4 paths; it is not a claim of universal multi-node or model support. Speculative steps are not included in mixed prefill/decode batches. + + ### Designed for recurrent inference - **Depth-aware paged KV cache:** LAST_EXITED and SHARED layouts, with optional @@ -92,6 +94,10 @@ in mixed prefill/decode batches. CUDA Graphs, and optional asynchronous/speculative paths. An optional native CUDA SiLU-and-Mul operator accelerates the Ouro activation path, with a Torch fallback; it is not a general model-speedup claim. +- **Sparse decode reference integration (opt-in):** cached loop-2 attention + correction for ordinary decode and block self-speculation, with shared + per-token reference ownership and explicit capability limits; see the + [configuration and scope](docs/runtime.md#sparse-decode-attention-reference-path). - **Offline and online use:** Python API, CLI, and a streaming OpenAI-compatible **completions** endpoint. This is not a chat-completions API. diff --git a/benchmarks/sparse_decode_attention_bench.py b/benchmarks/sparse_decode_attention_bench.py new file mode 100644 index 00000000..3f0fe63f --- /dev/null +++ b/benchmarks/sparse_decode_attention_bench.py @@ -0,0 +1,196 @@ +"""Sparse decode attention reference-path benchmark (RFC #156, section 3.2). + +Measures the reference implementation of the decode correction against the +dense reference path on identical paged-KV inputs: + +* ``dense`` one loop of full attention over ``N`` keys, +* ``capture`` the reference loop: dense attention plus retained statistics, +* ``consume`` a later loop: correction over ``k = ceil(N * retention)`` keys. + +A four-loop decode therefore costs ``4 x dense`` in the baseline and +``dense + capture + 2 x consume`` with sparse decode attention. This measures +the algorithm on the reference path, not an optimized kernel: the reported +speedup is a work proxy, not a serving claim (RFC #156 section 9). + +Usage: + python benchmarks/sparse_decode_attention_bench.py --device cuda + python benchmarks/sparse_decode_attention_bench.py --device cpu --context 512 +""" + +import argparse +import json +import math +import time + +import torch + +from vllm_rlt.attention.sparse import ( + reference_statistics, + retained_count, + sparse_correction, +) +from vllm_rlt.kernels.paged_attention import torch_paged_attention + + +def build_paged(context, heads, head_dim, kv_heads, block_size, device, dtype): + blocks = math.ceil(context / block_size) + shape = (blocks, block_size, kv_heads, head_dim) + generator = torch.Generator(device="cpu").manual_seed(0) + key_cache = torch.randn(shape, generator=generator, dtype=torch.float32).to(device, dtype) + value_cache = torch.randn(shape, generator=generator, dtype=torch.float32).to(device, dtype) + table = torch.arange(blocks, device=device, dtype=torch.int32).reshape(1, blocks) + lengths = torch.tensor([context], device=device, dtype=torch.int32) + q = torch.randn((1, heads, head_dim), generator=generator, dtype=torch.float32) + return key_cache, value_cache, table, lengths, q.to(device, dtype) + + +def timed(function, repeats, device): + """Median wall time of ``repeats`` calls, after warmup.""" + for _ in range(5): + function() + if device.type == "cuda": + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + function() + if device.type == "cuda": + torch.cuda.synchronize() + samples.append(time.perf_counter() - start) + samples.sort() + return samples[len(samples) // 2] * 1e3 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cpu") + parser.add_argument("--context", type=int, default=512) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--layers", type=int, default=24) + parser.add_argument("--loops", type=int, default=4) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"]) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + + device = torch.device(args.device) + dtype = getattr(torch, args.dtype) + key_cache, value_cache, table, lengths, q = build_paged( + args.context, args.heads, args.head_dim, args.kv_heads, args.block_size, device, dtype + ) + block_size = args.block_size + positions = torch.arange(args.context, device=device) + blocks = table[0, positions // block_size] + offsets = positions % block_size + keys = key_cache[blocks, offsets] + values = value_cache[blocks, offsets] + retained = retained_count(args.context, args.retention) + + dense_ms = timed( + lambda: torch_paged_attention(q, key_cache, value_cache, table, lengths), + args.repeats, + device, + ) + capture_ms = timed( + lambda: reference_statistics( + q[0], keys, values, key_retention=args.retention, key_bucket=retained + ), + args.repeats, + device, + ) + output, contribution, mass, selected = reference_statistics( + q[0], keys, values, key_retention=args.retention, key_bucket=retained + ) + groups = args.heads // args.kv_heads + retained_keys = key_cache[ + table[0, selected.long() // block_size], + selected.long() % block_size, + (torch.arange(args.heads, device=device) // groups)[:, None].expand_as(selected), + ] + retained_values = value_cache[ + table[0, selected.long() // block_size], + selected.long() % block_size, + (torch.arange(args.heads, device=device) // groups)[:, None].expand_as(selected), + ] + consume_ms = timed( + lambda: sparse_correction(q[0], retained_keys, retained_values, mass, output, contribution), + args.repeats, + device, + ) + corrected = sparse_correction(q[0], retained_keys, retained_values, mass, output, contribution) + # The reference loop and the later loop hold different KV, so the retained + # set chosen at loop 2 is applied to loop 3 keys. Compare against dense + # attention over that later KV to measure the approximation, not the + # same-loop identity. + later_keys = keys * 1.05 + 0.02 + later_values = values * 0.97 - 0.01 + later_retained_keys = retained_keys * 1.05 + 0.02 + later_retained_values = retained_values * 0.97 - 0.01 + dense_later = reference_statistics(q[0], later_keys, later_values, key_retention=1.0)[0] + dense_reference = reference_statistics(q[0], keys, values, key_retention=1.0)[0] + corrected_later = sparse_correction( + q[0], later_retained_keys, later_retained_values, mass, output, contribution + ) + divergence = float((corrected_later - dense_later).abs().max()) + same_loop_error = float((corrected - dense_reference).abs().max()) + + dense_total = args.layers * args.loops * dense_ms + sparse_total = args.layers * (dense_ms + capture_ms + (args.loops - 2) * consume_ms) + element = 2 if args.dtype == "bfloat16" else 4 + store_bytes = args.layers * args.heads * (retained * 4 + 4 + element * args.head_dim * 2) + kv_bytes = args.layers * 2 * args.kv_heads * args.head_dim * element + + report = dict( + device=str(device), + dtype=args.dtype, + context=args.context, + retained=retained, + retention=args.retention, + layers=args.layers, + heads=args.heads, + head_dim=args.head_dim, + dense_ms=round(dense_ms, 4), + capture_ms=round(capture_ms, 4), + consume_ms=round(consume_ms, 4), + dense_decode_ms=round(dense_total, 3), + sparse_decode_ms=round(sparse_total, 3), + decode_speedup=round(dense_total / sparse_total, 3), + keys_dense=args.context, + keys_sparse=retained, + key_ratio=round(retained / args.context, 4), + reference_state_bytes_per_slot=store_bytes, + kv_bytes_per_token=kv_bytes, + state_over_kv_ratio=round(store_bytes / kv_bytes, 3), + same_loop_identity_error=same_loop_error, + later_loop_max_abs_divergence=divergence, + ) + if args.json: + print(json.dumps(report, indent=2)) + return + print(f"sparse decode attention reference path — {report['device']} {args.dtype}") + print( + f" context={args.context} retained={retained} retention={args.retention} " + f"layers={args.layers} heads={args.heads} head_dim={args.head_dim}" + ) + print(f" dense row {dense_ms:8.4f} ms") + print(f" capture row {capture_ms:8.4f} ms") + print(f" consume row {consume_ms:8.4f} ms (k={retained})") + print( + f" {args.loops}-loop decode: dense {dense_total:8.3f} ms " + f"sparse {sparse_total:8.3f} ms speedup {dense_total / sparse_total:5.2f}x" + ) + print(f" keys touched per row: {args.context} -> {retained} ({retained / args.context:.1%})") + print( + f" reference state {store_bytes / 1024:.1f} KiB/slot vs KV " + f"{kv_bytes / 1024:.1f} KiB/token ({store_bytes / kv_bytes:.2f}x)" + ) + print(f" same-loop identity error (must be ~0): {same_loop_error:.3e}") + print(f" later-loop divergence at this retention: {divergence:.3e}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sparse_decode_batched_ab.py b/benchmarks/sparse_decode_batched_ab.py new file mode 100644 index 00000000..99541948 --- /dev/null +++ b/benchmarks/sparse_decode_batched_ab.py @@ -0,0 +1,169 @@ +"""A/B the scalar and batched sparse decode paths inside the merged tree. + +Both implementations live in ``vllm_rlt.attention.sparse`` and share the plan and +store contract, so this isolates the batching gain from the numerical contract: + +* ``scalar`` ``reference_paged_attention`` (row loop, the oracle), +* ``batched`` ``batched_paged_attention`` (equal-length row groups, CUDA). + +Capture rows are also measured with and without the dense-backend handoff, which +the full-retention control requires. + +Usage: + python benchmarks/sparse_decode_batched_ab.py --device cuda --rows 32 --context 512 +""" + +import argparse +import json +import math +import statistics +import time + +import torch + +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + batched_paged_attention, + reference_paged_attention, +) +from vllm_rlt.kernels.paged_attention import torch_paged_attention +from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + +def build(rows, context, heads, head_dim, kv_heads, layers, block_size, retention, device): + blocks = math.ceil(context / block_size) + generator = torch.Generator().manual_seed(0) + shape = (blocks, block_size, kv_heads, head_dim) + key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) + lengths = torch.full((rows,), context, device=device, dtype=torch.int32) + q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) + capacity = math.ceil(context * retention) + store = AttentionReferenceStore( + slots=rows, + layers=layers, + heads=heads, + head_dim=head_dim, + key_bucket=capacity, + device=device, + ) + return key, value, table, lengths, q, store + + +def timed(fn, repeats, device): + for _ in range(3): + fn() + torch.cuda.synchronize(device) + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize(device) + samples.append(time.perf_counter() - start) + return statistics.median(samples) * 1e3 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--rows", type=int, default=32) + parser.add_argument("--context", type=int, default=512) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--layers", type=int, default=24) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + device = torch.device(args.device) + + report = { + "device": str(device), + "rows": args.rows, + "context": args.context, + "retention": args.retention, + "layers": args.layers, + } + for mode, name in ((CAPTURE, "capture"), (CONSUME, "consume")): + key, value, table, lengths, q, store = build( + args.rows, + args.context, + args.heads, + args.head_dim, + args.kv_heads, + args.layers, + args.block_size, + args.retention, + device, + ) + lease = store.reserve("r", 0, args.rows, generation=1) + owners = [(lease, row) for row in range(args.rows)] + # Seed a capture pass so consume rows have state. + seed_plan = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) + reference_paged_attention( + q, + key, + value, + table, + lengths, + seed_plan, + layer=0, + dense_attention=torch_paged_attention, + ) + plan = store.build_plan(owners, [mode] * args.rows, key_retention=args.retention) + dense = torch_paged_attention + scalar = timed( + lambda: reference_paged_attention( + q, key, value, table, lengths, plan, layer=0, dense_attention=dense + ), + args.repeats, + device, + ) + batched = timed( + lambda: batched_paged_attention( + q, key, value, table, lengths, plan, layer=0, dense_attention=dense + ), + args.repeats, + device, + ) + batched_oracle = timed( + lambda: batched_paged_attention(q, key, value, table, lengths, plan, layer=0), + args.repeats, + device, + ) + report[name] = { + "scalar_ms": round(scalar, 3), + "batched_ms": round(batched, 3), + "batched_no_dense_handoff_ms": round(batched_oracle, 3), + "speedup": round(scalar / batched, 2), + } + store.release(lease) + + dense_ms = timed( + lambda: torch_paged_attention(q, key, value, table, lengths), args.repeats, device + ) + report["dense_attention_ms"] = round(dense_ms, 3) + if args.json: + print(json.dumps(report, indent=2)) + return + print( + f"sparse decode path A/B — {device}, rows={args.rows}, context={args.context}, " + f"retention={args.retention}, layers={args.layers}" + ) + for name in ("capture", "consume"): + values = report[name] + print( + f" {name:8s} scalar {values['scalar_ms']:7.3f} ms " + f"batched {values['batched_ms']:7.3f} ms " + f"speedup {values['speedup']:5.2f}x " + f"(batched w/o dense handoff {values['batched_no_dense_handoff_ms']:7.3f} ms)" + ) + print(f" dense attention reference: {dense_ms:.3f} ms") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sparse_decode_context_sweep.py b/benchmarks/sparse_decode_context_sweep.py new file mode 100644 index 00000000..6eae0687 --- /dev/null +++ b/benchmarks/sparse_decode_context_sweep.py @@ -0,0 +1,242 @@ +"""Dense-baseline context sweep and memory breakdown for sparse decode attention. + +The acceptance baseline is the dense Triton kernel, not the scalar sparse +reference: a four-loop decode costs ``4 x dense`` with the feature off and +``dense + capture + 2 x consume`` with it on. This script measures both on +identical paged KV at increasing context, reports the ratio, and breaks memory +down into KV storage, the persistent reference pool and transient workspace. + +Usage: + python benchmarks/sparse_decode_context_sweep.py --device cuda + python benchmarks/sparse_decode_context_sweep.py --contexts 512,2048,8192 --rows 32 +""" + +import argparse +import json +import math +import statistics +import time + +import torch + +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + reference_paged_attention, + sparse_paged_attention, +) +from vllm_rlt.kernels.paged_attention import triton_paged_attention +from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + +def build(context, rows, heads, head_dim, kv_heads, layers, block_size, retention, device): + blocks = math.ceil(context / block_size) + generator = torch.Generator().manual_seed(0) + shape = (blocks, block_size, kv_heads, head_dim) + key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) + lengths = torch.full((rows,), context, device=device, dtype=torch.int32) + q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) + capacity = math.ceil(context * retention) + store = AttentionReferenceStore( + slots=rows, + layers=layers, + heads=heads, + head_dim=head_dim, + key_bucket=capacity, + device=device, + ) + return key, value, table, lengths, q, store + + +def timed(fn, repeats, device): + for _ in range(2): + fn() + torch.cuda.synchronize(device) + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize(device) + samples.append(time.perf_counter() - start) + return statistics.median(samples) * 1e3 + + +def transient_bytes(fn, device): + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + before = torch.cuda.memory_allocated(device) + fn() + torch.cuda.synchronize(device) + return max(0, torch.cuda.max_memory_allocated(device) - before) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--contexts", default="256,512,1024,2048,4096,8192,16384,32768") + parser.add_argument("--rows", type=int, default=32) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--layers", type=int, default=24) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + device = torch.device(args.device) + contexts = [int(value) for value in args.contexts.split(",")] + + rows = [] + for context in contexts: + try: + key, value, table, lengths, q, store = build( + context, + args.rows, + args.heads, + args.head_dim, + args.kv_heads, + args.layers, + args.block_size, + args.retention, + device, + ) + lease = store.reserve("r", 0, args.rows, generation=1) + owners = [(lease, row) for row in range(args.rows)] + seed = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) + sparse_paged_attention( + q, key, value, table, lengths, seed, layer=0, dense_attention=triton_paged_attention + ) + capture_plan = store.build_plan( + owners, [CAPTURE] * args.rows, key_retention=args.retention + ) + consume_plan = store.build_plan( + owners, [CONSUME] * args.rows, key_retention=args.retention + ) + dense_ms = timed( + lambda: triton_paged_attention(q, key, value, table, lengths), + args.repeats, + device, + ) + capture_ms = timed( + lambda: sparse_paged_attention( + q, + key, + value, + table, + lengths, + capture_plan, + layer=0, + dense_attention=triton_paged_attention, + ), + args.repeats, + device, + ) + consume_ms = timed( + lambda: sparse_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + args.repeats, + device, + ) + scalar_consume_ms = timed( + lambda: reference_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + args.repeats, + device, + ) + dense_total = 4 * dense_ms + sparse_total = dense_ms + capture_ms + 2 * consume_ms + rows.append( + dict( + context=context, + retained=math.ceil(context * args.retention), + dense_ms=round(dense_ms, 3), + capture_ms=round(capture_ms, 3), + consume_ms=round(consume_ms, 3), + scalar_consume_ms=round(scalar_consume_ms, 3), + dense_4loop_ms=round(dense_total, 3), + sparse_4loop_ms=round(sparse_total, 3), + sparse_over_dense=round(sparse_total / dense_total, 3), + capture_transient_mib=round( + transient_bytes( + lambda: sparse_paged_attention( + q, + key, + value, + table, + lengths, + capture_plan, + layer=0, + dense_attention=triton_paged_attention, + ), + device, + ) + / 1024**2, + 1, + ), + consume_transient_mib=round( + transient_bytes( + lambda: sparse_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + device, + ) + / 1024**2, + 1, + ), + reference_pool_mib=round(rows_pool_bytes(args, context) / 1024**2, 1), + kv_mib=round(key.numel() * key.element_size() * 2 / 1024**2, 1), + ) + ) + store.release(lease) + torch.cuda.empty_cache() + except torch.OutOfMemoryError: + rows.append(dict(context=context, status="oom")) + torch.cuda.empty_cache() + + if args.json: + print(json.dumps(rows, indent=2)) + return + print( + f"dense-baseline sweep — {device}, rows={args.rows}, layers={args.layers}, " + f"heads={args.heads}, head_dim={args.head_dim}, retention={args.retention}" + ) + print( + f"{'ctx':>7}{'k':>7}{'dense':>9}{'capture':>9}{'consume':>9}" + f"{'4-loop dense':>13}{'4-loop sparse':>14}{'ratio':>8}" + ) + for row in rows: + if row.get("status") == "oom": + print(f"{row['context']:>7} OOM") + continue + print( + f"{row['context']:>7}{row['retained']:>7}{row['dense_ms']:>9.2f}" + f"{row['capture_ms']:>9.2f}{row['consume_ms']:>9.2f}" + f"{row['dense_4loop_ms']:>13.2f}{row['sparse_4loop_ms']:>14.2f}" + f"{row['sparse_over_dense']:>8.2f}x" + ) + print() + print(f"{'ctx':>7}{'capture ws':>12}{'consume ws':>12}{'ref pool':>11}{'KV':>9}") + for row in rows: + if row.get("status") == "oom": + continue + print( + f"{row['context']:>7}{row['capture_transient_mib']:>10.1f}Mi" + f"{row['consume_transient_mib']:>10.1f}Mi" + f"{row['reference_pool_mib']:>9.1f}Mi{row['kv_mib']:>7.1f}Mi" + ) + + +def rows_pool_bytes(args, context): + """Persistent reference pool for this sweep's row count.""" + capacity = math.ceil(context * args.retention) + per_slot = capacity * 4 + 4 + 4 * 2 * args.head_dim + 8 + return args.rows * args.layers * args.heads * per_slot + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sparse_decode_engine_e2e.py b/benchmarks/sparse_decode_engine_e2e.py new file mode 100644 index 00000000..edd330d6 --- /dev/null +++ b/benchmarks/sparse_decode_engine_e2e.py @@ -0,0 +1,156 @@ +"""Engine-level end-to-end comparison for sparse decode attention. + +Runs the real StageEngine with asynchronous self-speculation and compares dense +against sparse decode at full and reduced retention: throughput, peak memory and +token agreement. The per-layer sweeps say what one attention call costs; this +says whether that survives scheduling, sampling and delivery. + +The tiny Ouro architecture is used because no trained checkpoint is available on +this host, so the numbers are a scheduler-level check, not a serving result. +""" + +import argparse +import statistics +import time + +import torch + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import AutoModelForCausalLM + + +def build(model, concurrency, sparse, blocks=2048): + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(blocks, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=concurrency, + max_num_batched_tokens=8 * concurrency, + prefill_chunk_size=16, + ), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ) + + +def run(model, concurrency, sparse, *, prompt, max_tokens, repeats, blocks=2048): + """One configuration: returns median wall time, tokens/s and peak bytes.""" + outputs = None + walls = [] + peaks = [] + for _ in range(repeats): + engine = build(model, concurrency, sparse, blocks) + torch.cuda.reset_peak_memory_stats() + params = [ + SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) + ] + prompts = [ + [2 + (index + position) % 60 for position in range(prompt)] + for index in range(concurrency) + ] + start = time.perf_counter() + for index in range(concurrency): + engine.add_request(f"r{index}", prompts[index], params[index]) + collected = {} + for _ in range(100000): + for output in engine.step(): + collected[output.request_id] = output + if len(collected) == concurrency and all(o.finished for o in collected.values()): + break + torch.cuda.synchronize() + walls.append(time.perf_counter() - start) + peaks.append(torch.cuda.max_memory_allocated()) + outputs = {key: list(value.token_ids) for key, value in collected.items()} + tokens = concurrency * max_tokens + return { + "wall": statistics.median(walls), + "tokens_per_s": tokens / statistics.median(walls), + "peak_gib": max(peaks) / 2**30, + "outputs": outputs, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--prompt", type=int, default=32) + parser.add_argument("--max-tokens", type=int, default=16) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--concurrency", default="1,8,32") + parser.add_argument("--retentions", default="1.0,0.1") + parser.add_argument("--model-path", default=None) + parser.add_argument("--cache-blocks", type=int, default=2048) + parser.add_argument("--sparse-context", type=int, default=64) + args = parser.parse_args() + + device = args.device + if args.model_path: + model = AutoModelForCausalLM.from_pretrained( + args.model_path, device=device, dtype=torch.bfloat16 + ) + else: + model = tiny_ouro_model(device) + concurrency_levels = [int(value) for value in args.concurrency.split(",")] + retentions = [float(value) for value in args.retentions.split(",")] + + label = args.model_path or "tiny Ouro" + print( + f"engine e2e — {label} ({model.config.num_hidden_layers}L, ctx " + f"{model.config.max_position_embeddings}), prompt {args.prompt}, " + f"max_tokens {args.max_tokens}, median of {args.repeats}" + ) + header = f"{'conc':>5}{'case':>14}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}" + print(header) + for concurrency in concurrency_levels: + dense = run( + model, + concurrency, + None, + prompt=args.prompt, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.cache_blocks, + ) + print( + f"{concurrency:>5}{'dense':>14}{dense['tokens_per_s']:>10.2f}" + f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" + ) + for retention in retentions: + sparse = run( + model, + concurrency, + SparseDecodeAttentionConfig( + enabled=True, + key_retention=retention, + max_context_tokens=args.sparse_context, + ), + prompt=args.prompt, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.cache_blocks, + ) + same = sum( + 1 + for key, tokens in sparse["outputs"].items() + if tokens == dense["outputs"].get(key) + ) + print( + f"{'':>5}{f'sparse r={retention}':>14}{sparse['tokens_per_s']:>10.2f}" + f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" + f"{f'{same}/{concurrency}':>10}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sparse_decode_engine_e2e_text.py b/benchmarks/sparse_decode_engine_e2e_text.py new file mode 100644 index 00000000..734bbd9c --- /dev/null +++ b/benchmarks/sparse_decode_engine_e2e_text.py @@ -0,0 +1,159 @@ +"""Engine e2e with a real-text prompt, for a trained checkpoint. + +Synthetic token-id prompts drive a trained model into a fixed point (one token +repeated), which makes token agreement meaningless: every implementation agrees +when the argmax is constant. This encodes real text with the checkpoint's own +tokenizer and reports throughput, peak memory and per-request token agreement. + +Usage: + python benchmarks/sparse_decode_engine_e2e_text.py \ + --model-path /path/to/Ouro-1.4B --concurrency 1,8 +""" + +import argparse +import statistics +import time + +import torch +from tokenizers import Tokenizer + +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import AutoModelForCausalLM + +TEXT = ( + "The history of computing is a history of abstraction. Each generation of " + "hardware made a new kind of program possible, and each new kind of program " + "demanded a new way of thinking about machines. When memory was scarce, " + "programmers counted bytes by hand; when it became plentiful, they wrote " + "interpreters that treated memory as infinite. The same shift is happening " + "again with inference: models that once fit on a single device now span " + "racks, and the scheduler rather than the kernel decides how fast a request " + "finishes. Attention is the expensive part, because every token must look at " + "every earlier token, and that cost grows with the square of the context. " + "Sparse attention tries to keep only the keys that matter, which is easy to " + "state and hard to make fast, because deciding what matters is itself as " + "expensive as attending to everything." +) + + +def build(model, concurrency, sparse, blocks): + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(blocks, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=concurrency, + max_num_batched_tokens=8 * concurrency, + prefill_chunk_size=16, + ), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ) + + +def run(model, concurrency, sparse, *, prompt_ids, max_tokens, repeats, blocks): + walls = [] + peaks = [] + outputs = None + for _ in range(repeats): + engine = build(model, concurrency, sparse, blocks) + torch.cuda.reset_peak_memory_stats() + params = [ + SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) + ] + start = time.perf_counter() + for index in range(concurrency): + engine.add_request(f"r{index}", list(prompt_ids), params[index]) + collected = {} + for _ in range(100000): + for output in engine.step(): + collected[output.request_id] = output + if len(collected) == concurrency and all(o.finished for o in collected.values()): + break + torch.cuda.synchronize() + walls.append(time.perf_counter() - start) + peaks.append(torch.cuda.max_memory_allocated()) + outputs = {key: list(value.token_ids) for key, value in collected.items()} + tokens = concurrency * max_tokens + return { + "tokens_per_s": tokens / statistics.median(walls), + "wall": statistics.median(walls), + "peak_gib": max(peaks) / 2**30, + "outputs": outputs, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", required=True) + parser.add_argument("--prompt", type=int, default=200) + parser.add_argument("--max-tokens", type=int, default=16) + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument("--concurrency", default="1,8") + parser.add_argument("--blocks", type=int, default=768) + parser.add_argument("--sparse-context", type=int, default=512) + args = parser.parse_args() + + model = AutoModelForCausalLM.from_pretrained( + args.model_path, device="cuda", dtype=torch.bfloat16 + ) + tokenizer = Tokenizer.from_file(f"{args.model_path}/tokenizer.json") + ids = tokenizer.encode(TEXT).ids + while len(ids) < args.prompt: + ids = ids + ids + prompt_ids = ids[: args.prompt] + print(f"real text prompt: {args.prompt} tokens, head={prompt_ids[:8]}") + + print(f"{'conc':>5}{'case':>16}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}") + for concurrency in [int(v) for v in args.concurrency.split(",")]: + dense = run( + model, + concurrency, + None, + prompt_ids=prompt_ids, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.blocks, + ) + print( + f"{concurrency:>5}{'dense':>16}{dense['tokens_per_s']:>10.2f}" + f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" + ) + for retention in (1.0, 0.1): + sparse = run( + model, + concurrency, + SparseDecodeAttentionConfig( + enabled=True, + key_retention=retention, + max_context_tokens=args.sparse_context, + ), + prompt_ids=prompt_ids, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.blocks, + ) + same = sum( + 1 + for key, tokens in sparse["outputs"].items() + if tokens == dense["outputs"].get(key) + ) + print( + f"{'':>5}{f'sparse r={retention}':>16}{sparse['tokens_per_s']:>10.2f}" + f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" + f"{f'{same}/{concurrency}':>10}" + ) + print(f" dense sample tokens: {dense['outputs']['r0'][:8]}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ttft_bench.py b/benchmarks/ttft_bench.py new file mode 100644 index 00000000..9250710e --- /dev/null +++ b/benchmarks/ttft_bench.py @@ -0,0 +1,110 @@ +"""Engine-level TTFT / ITL / throughput comparison: dense vs sparse decode. + +Prefill rows never take the sparse decode correction (RFC #156 section 6.3), so +enabling the feature must not change prefill work. Any TTFT difference therefore +comes from the first speculative round, which does run the decode path. + +Usage: + python benchmarks/ttft_bench.py --device cuda --prompt 64 --max-tokens 32 + python benchmarks/ttft_bench.py --device cuda --retention 0.5 +""" + +import argparse +import json +import statistics +import time + +import torch + +from tests.helpers import make_speculative_engine, tiny_ouro_model +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig + + +def measure(engine, prompt, max_tokens, repeats): + """TTFT and inter-token latency of sequential single-request decodes.""" + ttfts, itls, totals = [], [], [] + for repetition in range(repeats): + request_id = f"r{repetition}" + start = time.perf_counter() + engine.add_request( + request_id, + prompt, + SamplingParams(max_tokens=max_tokens, ignore_eos=True), + ) + first = None + finished = None + while True: + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + if output.request_id != request_id: + continue + if first is None: + first = time.perf_counter() - start + if output.finished: + finished = time.perf_counter() - start + assert first is not None and finished is not None + ttfts.append(first * 1e3) + totals.append(finished * 1e3) + itls.append((finished - first) * 1e3 / max(1, max_tokens - 1)) + return dict( + ttft_median_ms=round(statistics.median(ttfts), 2), + ttft_min_ms=round(min(ttfts), 2), + itl_median_ms=round(statistics.median(itls), 3), + e2e_median_ms=round(statistics.median(totals), 2), + tokens_per_s=round(max_tokens / (statistics.median(totals) / 1e3), 1), + ) + + +def run(device, sparse, prompt, max_tokens, repeats, seed=123): + torch.manual_seed(seed) + model = tiny_ouro_model(device) + engine = make_speculative_engine(model, sparse=sparse) + # Warmup: compile-free but pays one-time kernel/allocator costs. + measure(engine, prompt, 4, 1) + result = measure(engine, prompt, max_tokens, repeats) + engine.close() + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--prompt", type=int, default=64) + parser.add_argument("--max-tokens", type=int, default=32) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + + prompt = [2 + (index % 60) for index in range(args.prompt)] + cases = { + "dense": None, + "sparse r=1.0": SparseDecodeAttentionConfig(enabled=True, key_retention=1.0), + "sparse r=0.5": SparseDecodeAttentionConfig(enabled=True, key_retention=0.5), + } + report = {} + for name, sparse in cases.items(): + report[name] = run(args.device, sparse, prompt, args.max_tokens, args.repeats) + if args.json: + print(json.dumps(report, indent=2)) + return + header = f"{'case':<14}{'TTFT med':>10}{'TTFT min':>10}{'ITL':>9}{'e2e':>9}{'tok/s':>9}" + print( + f"TTFT / ITL — {args.device}, prompt {args.prompt}, max_tokens {args.max_tokens}, " + f"repeats {args.repeats}" + ) + print(header) + for name, values in report.items(): + print( + f"{name:<14}{values['ttft_median_ms']:>9.2f}ms{values['ttft_min_ms']:>9.2f}ms" + f"{values['itl_median_ms']:>8.3f}ms{values['e2e_median_ms']:>8.1f}ms" + f"{values['tokens_per_s']:>9.1f}" + ) + base = report["dense"]["ttft_median_ms"] + for name in list(report)[1:]: + delta = report[name]["ttft_median_ms"] / base + print(f" TTFT {name} vs dense: {delta:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/docs/runtime.md b/docs/runtime.md index 460c5e9c..3869d9f0 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -230,6 +230,83 @@ The same flags work with the single-device server. A streaming event may contain multiple committed tokens. See [RFC #43](https://github.com/InferMatrixLab/vllm-rlt/issues/43) for the design. +## Sparse decode attention (reference path) + +`SparseDecodeAttentionConfig` enables the RFC #156 cached-reference correction +for generated-token attention. Prompt prefill remains dense. Loop 2 captures +per-token/per-layer/head selection, mass and contributions; loops 3 and 4 reuse +that reference with their current selected KV. Each new input token establishes +its own reference. + +```python +from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig, SpeculativeConfig + +sparse = SparseDecodeAttentionConfig( + enabled=True, key_retention=0.1, max_context_tokens=512, +) +# Pass sparse_decode_attention_config=sparse to LLM or StageEngine. +# Ordinary decode uses no speculative_config, with synchronous or +# asynchronous scheduling. For block speculation, also pass: +speculation = SpeculativeConfig(num_speculative_tokens=2, draft_loops=2) +execution = ExecutionConfig(async_scheduling=False, cuda_graphs=False) +# Token pipelines, and CUDA Graphs with speculation, need the fused kernels: +fused = SparseDecodeAttentionConfig(enabled=True, key_retention=0.1, implementation="fused") +pipeline = SpeculativeConfig(num_speculative_tokens=3, draft_loops=2, mode="pipeline") +execution = ExecutionConfig(async_scheduling=True, cuda_graphs=True) +``` + +This path runs with LAST_EXITED KV and refill scheduling in these modes: + +| Mode | Implementations | +| --- | --- | +| Ordinary decode, synchronous or asynchronous scheduling | `batched`, `fused` | +| Block self-speculation, synchronous or asynchronous | `batched`, `fused` | +| Asynchronous block self-speculation with CUDA Graphs | `fused`, `key_retention < 1` | +| Asynchronous token pipelines, with or without CUDA Graphs | `fused`, `key_retention < 1` | + +The Torch and Triton adapters declare sparse-reference support; `auto` selects +a compatible adapter and explicit FlashAttention selection fails rather than +silently replacing the requested backend. Dense and capture rows preserve the +selected backend's output and original packed batch geometry. In the default +`batched` implementation, selection, retained mass and contributions, and +partial-retention correction use the shared PyTorch reference implementation; +it has host readbacks and per-row launch costs. The opt-in `fused` +implementation (`vllm_rlt/kernels/sparse_decode.py`) keeps lengths and counts +on the device and launches a fixed set of Triton kernels per layer, so it can +be replayed inside CUDA Graphs; its consume rows differ from the reference in +FP32 reduction order. The captured output is the value actually entering the +recurrence, stored in FP32 after the backend's output-dtype conversion. +Enabling sparse decode does not by itself promise faster serving. Token +pipelines decide each row's depth on the device, so their plans carry +per-row modes and slots on the device: every request owns one reference slot +per pipeline row for its lifetime, and the slots move with the rows. Mixed +scheduling, preemption/recovery, the synchronous token pipeline and SHARED +KV are rejected explicitly. + +The retained count is `min(N, max(1, ceil(N * key_retention)))`, including the +current position among the N causal keys. `key_bucket` controls buffer capacity, +not the algorithm: it must hold the retained count at the supported context +bound. `max_context_tokens=None` uses the model context; a smaller explicit +serving bound reduces pool memory and is enforced before request admission. +The pool holds one slot per live ordinary input token, K+1 per synchronous +block request, two rounds of K+1 per asynchronous block request, and one per +pipeline row (at most K+1) per token-pipeline request. All +reference payload bytes contribute to automatic KV memory planning; explicit +KV block overrides still require the caller to budget total peak memory. + +Ordinary references retire at CODA or request teardown after their last +recurrent reader. Synchronous block references retire on the compute stream +when every verification input has completed its final read; neither accepted +nor rejected rows carry a reference into the next round. With draft depth 1, +verification captures loop 2; with draft depth 2, drafting captures it. +Full retention uses the same dense backend at later loops as a control, while +still validating and consuming the captured reference lifecycle. The standalone +Torch arithmetic oracle is tested separately; different attention algorithms +need not round to identical BF16 values. See the +[numerical diagnosis and regression report](validation/sparse-decode-numerics-20261029.md). +Retention below 1 is approximate and +needs task-quality evaluation before a deployment recommendation. + ## SHARED prefill semantics SHARED retains one physical KV plane and completes every loop for a prompt diff --git a/docs/validation/sparse-decode-numerics-20261029.md b/docs/validation/sparse-decode-numerics-20261029.md new file mode 100644 index 00000000..106918c7 --- /dev/null +++ b/docs/validation/sparse-decode-numerics-20261029.md @@ -0,0 +1,145 @@ +# Sparse decode: backend numerical diagnosis + +This follow-up fixes the full-retention discrepancy in the +[ordinary/synchronous integration report](sparse-decode-sync-20261029.md), within +the same Draft collaboration PR targeting `kunkun/rfc156-sparse-decode-async`. +The immediate feature base remains `048b9d7c79a5711857ce3287812502ad8f274ad0`; +the pre-fix contribution head is `7f22af3901416096ea6c1c0e0a5536f440f662f5`. +The fixed runtime/test commit is `1bbe6fdd24be3e9bbea6a9dd1440d4b64f8dd89b`. + +## Reproduced cause + +The original BF16 Ouro-1.4B workload reproduces 498/512 token agreement at +concurrent arrivals 32 and synchronous draft depth 1. Request 9 first differs +at output index 2: dense chooses 293, the original sparse reference chooses 92. +The 14 differences include subsequent autoregressive continuation, rather than +14 independent initial errors. + +The new diagnostic retains the original 128-input/16-output protocol, including +four-token warmup. It freezes the complete accepted verification batch's input +hidden states, logical row/depth/position metadata, position IDs and context +lengths, all physical-layer module outputs, full packed queries and attention +outputs, and the watched row's causal K/V. It filters out rejected executions of +the same position. The accepted loop-2 batch here has **22 rows**, not the +theoretical 96-row maximum; both arms have exactly equal packed input states, +request histories, shapes, positions and lengths. A shorter-output exploratory +run changed the batch geometry and is retained separately, not substituted for +this reproduction. + +For the watched row, physical layers 0 and 1 agree. At physical layer **2** +(zero-based), Q, all 130 legal keys and values are identical, but the attention +output differs at one BF16 element by **0.000030517578125**. The next `o_proj` +differs at two elements, max **0.000244140625**. Across the complete batch, +attention differences already exist at physical layer 0. Later layers and +recurrences amplify these differences. This places the first observed watched +operator discrepancy at attention, before reference consumption. + +A separate replay uses the frozen operands with no model or scheduler. Triton +split-KV reproduces the captured dense output exactly; the Torch +einsum/softmax reference reproduces the captured sparse output exactly. Changing +only the backend reproduces the layer-2 difference. Disabling BF16 reduced +precision reduction does not change that reference result; its reductions +already use FP32 operands. TF32 is disabled throughout. Equal packed shapes +rule out a batch-shape change as the explanation for this captured attention +discrepancy. No global precision flags or model GEMM implementations are changed. + +## Fix and numerical contract + +`KVCacheManager` passes its selected dense attention adapter into the shared +`sparse_paged_attention` dispatcher. The adapter receives the original packed +queries, tables and lengths, without slicing into different batch shapes. + +- DENSE and CAPTURE rows use that backend's actual output. Capture stores that + output as FP32 **after** the backend's output-dtype rounding, so `O_ref` is + the value entering recurrence. Retained selection, mass and contribution + remain the FP32 Torch reference arithmetic. +- At retention 1, CONSUME rows also use the configured dense backend, while + still checking reference ownership/payload availability and recording + consumption. This is explicitly a dense control; it does not qualify the + Torch correction as bitwise equal to Triton. +- At retention below 1, later loops retain `O_ref - C_ref + m * C_new`. This + changes the original reference's source-output rounding contract, not the + selection or sparse delta formula. Approximate quality remains unqualified. +- Omitting the adapter when calling the standalone paged oracle forces the + original Torch capture and correction, including at retention 1. Known-answer + paged cases test changed later keys/values at retention .5 and 1; the full-retention + control cannot hide a regression in that independent arithmetic path. + +A full-retention-only bypass would leave partial-retention capture rows using +a different dense algorithm; preserving the selected backend at every dense +source row fixes that inconsistency too. No new kernel or global precision +setting is introduced. The fallback oracle remains independently available. + +## Regression evidence + +Four CUDA BF16 cases cover MHA/GQA, retention .1/1, a mixed dense/capture plan, +stored `O_ref`, changed later-loop operands and reference consumption. All four +fail on the original implementation with backend-output differences, then pass +with the fix using exact equality. The standalone paged known-answer cases +return 23.5 and 40.0 after changing later keys/values, exercising partial and +full Torch correction without the backend control. + +Full CPU regression: **921 passed, 479 skipped**. Full pre-commit passes Ruff +lint, Ruff format and the CPU suite. Independent read-only review found no +actionable defect in the four-file runtime/test delta; its targeted CPU check +passed 48 cases, with 22 GPU cases deselected. + +Applicable CUDA regression: **318 passed, 1 skipped, 16 deselected**. The skip +requires two GPUs; FA4 cases remain excluded on this SM89 device. A fixed +instrumented replay now matches the original dense replay exactly at every +recorded depth: complete packed inputs/outputs, all physical module outputs, +attention queries/outputs and the watched row's causal K/V. This comparison is +separate from the uninstrumented generation matrix. + +The uninstrumented BF16 regression matrix covers ordinary decode and synchronous +draft depths 1/2, simultaneous arrivals 1/8/32, and retention off/1/.1. Each of +the 27 fresh-process configurations warms four output tokens, then measures +three bursts using the original 128-input/16-output protocol. All **81 bursts, +1,107 requests and 17,712 output tokens** complete, with KV/reference state +drained. Tokens and exit depths repeat exactly within each configuration. + +| Mode | C=1 full vs dense | C=8 full vs dense | C=32 full vs dense | +|---|---:|---:|---:| +| Ordinary | 16/16 | 128/128 | 512/512 | +| Sync, draft 1 | 16/16 | 128/128 | 512/512 | +| Sync, draft 2 | 16/16 | 128/128 | 512/512 | + +Every cell matches in all three repetitions, including exit depths: **5,904/5,904 +full-retention tokens** in total. In particular, the original C32/draft-1 case +improves from 498/512 to 512/512 in each repetition. This verifies the explicit +dense-backend control; the forced Torch correction remains a separate oracle. +Retention .1 completes and drains, but its token agreement is not a quality score. + +## Reproduction and archive + +The [numerical validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029) +contains the frozen tensor snapshots, source and SHA256 inventories, Git bundle, +same-input replay, original failures, red/green regression logs, full raw matrix, +environment metadata and reproduction commands. Model weights are not included. +The original unsupported-FA4 failures and performance profiles remain in the +[earlier release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029). + +Use `trace_first_operator.py` against `7f22af3` for the dense and pre-fix capture, +then against `1bbe6fd` for the fixed capture; `replay_attention.py` tests the +frozen operands independently. `compare_fixed_trace.py` checks every recorded +state, and `run_regression_matrix.py` / `summarize_matrix.py` reproduce and verify +the complete generation matrix. These local experiment scripts preserve their +recorded root directory; the archive README explains relocation and exact CLI +commands. Diagnostic timings are excluded from uninstrumented results. + +## Scope + +The model weights match the prior SHA256 +`58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. +The follow-up server reports one RTX4090 (SM89) with 49,140 MiB, driver +595.71.05, torch 2.12.1+cu130 and Triton 3.7.1. This is a different instance +from the original 24,564-MiB server with driver 580.105.08; +timings across these instances are not a matched performance comparison. + +These are numerical and execution regressions on synthetic fixed-token inputs, +not a held-out task-quality evaluation. Partial retention, adaptive exits, +long contexts, other GPUs and untested feature combinations remain outside the +qualification. Existing asynchronous lifecycle regressions are included, but +no new real-model asynchronous quality claim is made. The original performance +profiles remain historical evidence; +this correctness fix makes no serving-speed improvement claim. diff --git a/docs/validation/sparse-decode-sync-20261029.md b/docs/validation/sparse-decode-sync-20261029.md new file mode 100644 index 00000000..91d3a7a1 --- /dev/null +++ b/docs/validation/sparse-decode-sync-20261029.md @@ -0,0 +1,96 @@ +# Sparse decode: ordinary and synchronous integration validation + +This is the historical report for runtime `10373ab`. The subsequent +[numerical diagnosis and regression report](sparse-decode-numerics-20261029.md) +documents the backend-preserving fix and its separate validation. The timings +and numerical limitation below describe the original implementation. + +This contribution extends the shared [#157](https://github.com/InferMatrixLab/vllm-rlt/pull/157) reference plan/store into ordinary synchronous decode and synchronous block self-speculation. It is a functional reference integration, with measured costs and an unresolved BF16 full-retention sequence difference. It does not claim an optimized sparse kernel or deployment quality. + +## Source and checks + +- Immediate base: `048b9d7c79a5711857ce3287812502ad8f274ad0` on `kunkun/rfc156-sparse-decode-async`. +- Measured runtime: `10373ab6d63156535b55e4b9be1c7b8ba7319421`. Additional same-input tests: `985c309`; runtime files are identical between these revisions. Source SHA256 inventories are in the archive. +- Full pre-commit: Ruff lint, Ruff format and CPU regression passed; **919 passed, 475 skipped**. +- Applicable GPU regression: **312 passed, 1 skipped, 16 deselected**. The skip needs a second GPU; FA4 cases are excluded on SM89. The original unfiltered run's 12 unsupported-FA4 failures are retained. +- New integration module: 26 CPU tests plus 5 CUDA cases. Includes ordinary/padded routing, prompt exclusion, draft depth 1/2, each rejection boundary, same-position replacement, abort/ID reuse, failed submissions, stream-owned retirement and active-request close. Per-depth hidden, logits and populated KV match serial same-input sparse execution at FP32 atol/rtol `4e-5`. +- The negative stream control restores the unsafe caller-stream release only inside one test process and fails as expected. Corrected production code passes. +- Independent review found release-index stream ownership and shutdown-drain defects; both were fixed and covered. Re-review found no further concrete production defect. +- Upstream FlashLoop `aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6`, `dense_source_selector` plus `cached_mass_sparse_delta(force_torch=True)`: **18 CUDA fixed-input cases passed**, with equal retained sets; lengths 1/13/128, retention .1/.5/1, FP32 and BF16. Predeclared attention-output atol/rtol were `2e-5` and `.02`, respectively, accommodating the gold path's BF16 subset/output storage. This is separate from model-level quality. + +## Real-model protocol + +One RTX4090, 24,564 MiB; driver 580.105.08; torch 2.12.1+cu130, Triton 3.7.1. Released Ouro-1.4B revision `574fa66cb8bf5abdc979642d01cf2b79b16bfab1`, BF16 weights/activations/KV; weight SHA256 `58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. + +Native StageEngine, LAST_EXITED KV, refill, Triton dense backend, synchronous scheduling, multistream on, graphs off. Block K=2, draft depths 1/2, fixed four loops, greedy and ignore EOS. Each fresh process warms 4 output tokens at its target arrival concurrency, then measures three 128-input/16-output bursts. Request i uses `[1 + (j * 137 + i * 53) % 49000 for j in range(128)]`. Context bound is explicitly **512** for sparse, while model positional configuration remains unchanged. Token budget `max(128, 3*C)`, prefill chunk 128, KV blocks `40*C + 32`, block size 16. This does not qualify the model's 65,536 context limit. + +Timing includes enqueue, prefill, decode, delivery, drain and CUDA synchronization; model loading and warmup are excluded. Profiles are separate and excluded. GPU clocks and CPU affinity were not pinned for the primary sweep. All raw repeats, outputs, depths, acceptance counters, TTFT/E2E observations and memory values are archived. Speculative output arrives in chunks; chunk gaps are not per-token ITL. The short bursts do not establish steady-state HTTP serving. + +**81 primary bursts / 1,107 completed requests / 17,712 output tokens**; no primary OOM, timeout or failed request. Every measured run drained KV blocks and reference slots to zero and repeated its token/depth sequences exactly. + +Output throughput: median of three runs, with min-max in parentheses. Full-retention agreement counts one repeat, and was identical across all three. + +| Mode | Concurrent arrivals | Dense tok/s | Full retention tok/s | r=.1 tok/s | Full vs dense tokens | +|---|---:|---:|---:|---:|---:| +| Ordinary | 1 | 14.90 (14.67-14.98) | 9.66 (8.89-9.79) | 9.30 (8.86-9.37) | 16/16 | +| Ordinary | 8 | 71.60 (71.59-71.60) | 21.59 (21.53-21.62) | 21.73 (21.59-21.77) | 128/128 | +| Ordinary | 32 | 124.42 (124.41-124.49) | 27.13 (27.01-27.31) | 26.84 (26.52-27.24) | 512/512 | +| Sync, draft 1 | 1 | 17.46 (17.21-17.68) | 9.10 (8.63-9.12) | 13.90 (12.72-14.01) | 16/16 | +| Sync, draft 1 | 8 | 60.53 (57.33-61.12) | 18.01 (17.83-19.32) | 25.54 (25.50-25.64) | 128/128 | +| Sync, draft 1 | 32 | 82.91 (81.09-83.10) | 19.58 (19.54-19.81) | 27.48 (27.05-27.57) | 498/512 | +| Sync, draft 2 | 1 | 16.56 (15.37-16.69) | 9.49 (9.06-9.70) | 11.87 (11.13-11.91) | 16/16 | +| Sync, draft 2 | 8 | 50.27 (50.16-50.57) | 18.55 (17.54-18.58) | 22.40 (22.23-22.42) | 128/128 | +| Sync, draft 2 | 32 | 68.08 (67.59-68.30) | 20.31 (20.20-20.38) | 24.84 (24.70-24.90) | 512/512 | + +The shared PyTorch reference path is slower than dense Triton in every covered condition. Peak allocated memory reached **6.67 GiB**. Retention .1 changes outputs, and its speculation acceptance changes are preserved in `summary.json`; token agreement is not a task-quality score. + +## Full-retention discrepancy + +At C=32, draft depth 1, full retention matches **498/512** dense output tokens. The first divergence is request 9, output index 2 (zero-based), input position 129: dense chooses 293, full retention chooses 92. Other full-retention conditions in this sweep match all tokens; this does not prove broader equivalence. + +Instrumented replays reproduce the mismatch with **identical 130-token input prefixes**. Loop-1 hidden and loop-2 input hidden are exact. The first observed hidden drift is already at loop-2 capture output (max .01953125), before any reference consume; later-loop input differences follow. Recorded layer-0 Q/K/V and attention are exact at the capture loop; their same-input full-retention attention replays match Torch. Final logits differ by at most .125. Dense ranks `(293, 8.875)` above `(92, 8.8125)`; the reference path reverses those values. This localizes the observed drift to the capture computation before reuse, consistent with backend numerical differences, but does not localize the first physical-layer operator or establish task quality. No lossless/full-retention-at-scale claim is made. The existing async discrepancy in #157 is also not resolved by this contribution. + +## Disabled-feature drift control + +Initial unpinned C8 base/candidate controls showed timing drift, so a separate CPU-pinned **ABBA** check was run for sync draft 2, C8, feature off. Each arm uses a fresh process, the same workload and three measurements; CPUs are `[0, 1, 2, 3]`. Tokens and exit depths match the base exactly in all arms. + +- Base: median 52.06, range 51.90-53.91 tok/s. +- Candidate: median 52.19, range 50.63-52.54 tok/s. +- Candidate/base median ratio: 1.0027. These pinned observations are a separate control, not replacements for the primary unpinned numbers or a general performance guarantee. + +## Separate profiling + +Twelve readable CPU+CUDA captures cover ordinary C1/C8 and sync draft-2 C8, dense versus r=.1, each in independent ops-only/full processes. Full enables shapes, memory and stacks together. Every full capture has shape events, memory events and a nonempty exported stack inventory. The profile window spans a further output from each active request after the first outputs; raw shapes and row work remain available for checking scheduling/acceptance differences. Diagnostic count ratios are not throughput ratios. + +| C8 arrival workload | Active requests dense / r=.1 | GPU kernels dense / r=.1 | Scalar item reads dense / r=.1 | Argsorts dense / r=.1 | +|---|---:|---:|---:|---:| +| Ordinary | 8 / 8 | 5,600 / 26,220 | 0 / 384 | 0 / 192 | +| Sync, draft 2 | 7 / 5 | 11,319 / 43,242 | 0 / 720 | 0 / 360 | + +Ordinary C8 windows contain eight active requests in both arms. The synchronous windows contain seven versus five active requests because progress differs; they support operator diagnosis, not a matched-work count ratio. + +The source chain is `reference_plan -> sparse_paged_attention -> reference_statistics/read_reference`: per-row gather, selection and scalar count reads add launches and waits. Pool operations also add retirement work. The measured slowdown is therefore not explained by retained-key count alone. This contribution retains the common reference oracle; it does not duplicate or replace the pending attention optimization discussed on #157. + +![Aligned single-request diagnostic decode windows](https://github.com/Levius-Fubuki/vllm-rlt/releases/download/pr157-sync-validation-20261029/single-request-timeline.png) + +## Reproduction and evidence + +The [validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029) contains scripts, original traces, stack exports, raw outputs, source inventories, the negative control, failed diagnostic-harness attempts and the first unsupported-FA4 run. The SHA256 inventory covers the package contents; `SHA256SUMS` covers the downloadable archive and timeline. + +From the corresponding source checkout with the recorded dependencies and a reserved GPU: + +```bash +pre-commit run --all-files +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python -m pytest \ + tests/test_sparse_decode_sync.py tests/test_sparse_decode_attention.py \ + tests/test_speculative.py tests/test_async_speculative.py \ + tests/test_async_speculative_lifecycle.py tests/test_backend_selection.py \ + --run-gpu -k 'not flash_attn_4' -q +PYTHONPATH=. CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 python scripts/real_model_bench.py \ + --model /path/to/pinned/Ouro-1.4B --mode sync2 --concurrency 32 \ + --retention .1 --repeats 3 --output /tmp/sync2-c32-r01.json +``` + +Use `--mode ordinary|sync1|sync2`, `--retention 0|1|.1`, `--concurrency 1|8|32`; retention 0 disables the feature. The script contains the exact engine startup and per-request parameters. Use `--profile ops|full --repeats 1` only for separate diagnostics. The orchestrators retain their original task directory for exact provenance; the individual benchmark CLI accepts arbitrary model/output locations. + +Graphs, token pipelines, mixed scheduling, preemption and ordinary async decode are explicitly rejected. PD, prefix-cache composition, adaptive-exit quality, long generations, held-out task quality and other hardware are not qualified here. Sparse remains disabled by default. This is a Draft collaboration contribution to the existing shared feature PR, with numerical/performance limits retained for review. diff --git a/tests/helpers.py b/tests/helpers.py index cf4a4236..ba4f3064 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -146,7 +146,7 @@ def prefill_submitted(self, request, **submission): def make_speculative_engine( - model, *, asynchronous=True, k=3, budget=16, prefix=False, multi_stream=True + model, *, asynchronous=True, k=3, budget=16, prefix=False, multi_stream=True, sparse=None ): return StageEngine( model, @@ -157,6 +157,7 @@ def make_speculative_engine( ), execution_config=ExecutionConfig(async_scheduling=asynchronous, multi_stream=multi_stream), speculative_config=SpeculativeConfig(k), + sparse_decode_attention_config=sparse, ) diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index e6179d0d..57735bcc 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -53,6 +53,30 @@ def test_cpu_auto_reports_torch_and_runs_reference_attention(caplog): torch.testing.assert_close(selected(q, k, v, table, lengths), explicit(q, k, v, table, lengths)) +def test_sparse_auto_filters_flash_and_explicit_flash_fails(packages, monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0)) + required = BackendCapabilities(sparse_decode=True) + selected = create_backend( + "auto", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + assert selected.info["selected_backend"] == "triton" + assert "sparse_decode" in selected.info["selection_reason"] + with pytest.raises(ValueError, match="lacks required capabilities: sparse_decode"): + create_backend( + "flash_attn_4", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + + @pytest.mark.parametrize( "sm,dtype,head_dim,block_size,missing,expected", [ diff --git a/tests/test_pd.py b/tests/test_pd.py index 90a4b38f..f1f553c4 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -13,6 +13,7 @@ ExitConfig, SamplingParams, SchedulerConfig, + SparseDecodeAttentionConfig, SpeculativeConfig, ) from vllm_rlt.core.kv_cache_manager import KVCacheManager @@ -118,6 +119,47 @@ def test_pd_rejects_huginn_before_spawning(monkeypatch): clients.assert_not_called() +def test_pd_forwards_sparse_decode_only_to_decode_workers(monkeypatch): + from unittest.mock import Mock + + from vllm_rlt.pd.engine import PDEngine + + config = tiny_ouro_config(head_dim=64) + sparse = SparseDecodeAttentionConfig(enabled=True, key_retention=0.5) + + def ready_client(model, *, stage_config, **options): + return Mock( + ready=True, + connected=True, + stopped=True, + info=dict( + kind="ready", + stage_id=stage_config.stage_id, + role=stage_config.role, + fingerprint="same-model-and-cache", + info={"agent": stage_config.stage_id}, + num_blocks=64, + block_size=16, + depths=model.total_ut_steps, + ), + ) + + clients = Mock(side_effect=ready_client) + monkeypatch.setattr("vllm_rlt.pd.engine.StageClient", clients) + with PDEngine( + config, + attention_backend="triton", + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ): + roles = { + call.kwargs["stage_config"].role: call.kwargs["sparse_decode_attention_config"] + for call in clients.call_args_list + } + # Prefill rows are always dense: only the decode stage owns reference state. + assert roles == {"prefill": None, "decode": sparse} + + def test_pd_nanbeige_startup_uses_entry_config_for_request_validation(monkeypatch): from unittest.mock import Mock @@ -317,6 +359,9 @@ def create_pd( speculative=False, asynchronous=True, prefix_caching=False, + sparse=None, + decode_batched_tokens=2, + speculative_mode="block", ): pytest.importorskip("nixl") from vllm_rlt.pd.engine import PDEngine @@ -343,11 +388,14 @@ def create_pd( prefill_chunk_size=prefill_chunk_size, wavefront_chunked_prefill=wavefront_chunked_prefill, ), - decode_scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=2), + decode_scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=decode_batched_tokens + ), exit_config=ExitConfig("ouro") if speculative else ExitConfig("trace", depths_by_request={"t": [4, 2, 3, 1] * 16}), - speculative_config=SpeculativeConfig(3) if speculative else None, + speculative_config=SpeculativeConfig(3, mode=speculative_mode) if speculative else None, + sparse_decode_attention_config=sparse, execution_config=ExecutionConfig( async_scheduling=asynchronous, static_buffers=graph, @@ -682,6 +730,103 @@ def test_pd_speculative_matches_local_and_retires_cancelled_rounds( assert all(metrics["used_blocks"] == 0 for metrics in engine.worker_metrics.values()) +@pytest.mark.gpu +@pytest.mark.parametrize( + "sparse", + [ + SparseDecodeAttentionConfig(enabled=True, key_retention=1.0), + SparseDecodeAttentionConfig(enabled=True, key_retention=0.5), + SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused"), + ], +) +def test_pd_sparse_decode_on_the_decode_stage(sparse): + params = SamplingParams(max_tokens=9, ignore_eos=True) + prompts = [[1, 2, 3, 4, 5], [2, 3], [4], [5, 6, 7]] + with create_pd(speculative=True, asynchronous=True, sparse=sparse) as engine: + for index, prompt in enumerate(prompts): + engine.add_request(str(index), prompt, params) + outputs = drain(engine) + assert sorted(outputs) == ["0", "1", "2", "3"] + assert all(len(output.token_ids) == 9 for output in outputs.values()) + if sparse.key_retention == 1.0: + # Full retention is the exact control: it must reproduce local dense decoding. + torch.manual_seed(123) + reference = LLMEngine( + OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda:0", torch.bfloat16), + cache_config=CacheConfig(128, 4), + attention_backend="triton", + ) + try: + for index, prompt in enumerate(prompts): + reference.add_request(str(index), prompt, params) + assert outputs == drain(reference) + finally: + reference.close() + # A request aborted mid-decode must still release its reference slots. + engine.add_request("abort", prompts[0], replace(params, max_tokens=60)) + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if any(len(output.token_ids) > 1 for output in engine.step()): + break + else: + pytest.fail("sparse speculative decode did not start") + engine.abort_request("abort") + drain(engine) + assert all(peer.slots == 0 and peer.reserved_blocks == 0 for peer in engine.peers.values()) + assert all(metrics["used_blocks"] == 0 for metrics in engine.worker_metrics.values()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("mode", ["block", "pipeline"]) +@pytest.mark.parametrize("multi", [False, True]) +def test_pd_sparse_decode_replays_graphs_on_the_decode_stage(multi, mode): + fused = SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused") + params = SamplingParams(max_tokens=9, ignore_eos=True) + prompts = [[1, 2, 3, 4, 5], [2, 3], [4], [5, 6, 7]] + results = {} + for graph, sparse in [(False, None), (True, None), (False, fused), (True, fused)]: + # Room for full verification rounds, so consume steps replay graphs too. + with create_pd( + graph=graph, + multi=multi, + speculative=True, + asynchronous=True, + sparse=sparse, + decode_batched_tokens=8, + speculative_mode=mode, + ) as engine: + for index, prompt in enumerate(prompts): + engine.add_request(str(index), prompt, params) + outputs = drain(engine) + assert sorted(outputs) == ["0", "1", "2", "3"] + results[graph, sparse is not None] = outputs + engine.add_request("abort", prompts[0], replace(params, max_tokens=60)) + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + if any(len(output.token_ids) > 1 for output in engine.step()): + break + else: + pytest.fail("speculative decode did not start") + engine.abort_request("abort") + drain(engine) + assert all( + peer.slots == 0 and peer.reserved_blocks == 0 for peer in engine.peers.values() + ) + metrics = engine.worker_metrics + assert all(m["used_blocks"] == 0 for m in metrics.values()) + decode = [m for name, m in metrics.items() if "-decode-" in name] + assert len(decode) == (2 if multi else 1) + if graph: + assert all(m["graph_replays"] > 0 for m in decode) + # Every decode worker owns its store and replays its sparse graphs. + assert all( + (m["sparse_graph_replays"] > 0) == (graph and sparse is not None) for m in decode + ) + assert results[True, False] == results[False, False] + # The replayed fused kernels are the eager ones, row for row. + assert results[True, True] == results[False, True] + + def test_pd_reservations_are_reported_before_worker_acknowledgments(): from types import SimpleNamespace diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py new file mode 100644 index 00000000..b21ca33b --- /dev/null +++ b/tests/test_sparse_decode_attention.py @@ -0,0 +1,811 @@ +"""Sparse decode attention: reference lifecycle and plan contracts (RFC #156). + +Covers the asynchronous self-speculation cell of the sparse decode attention +column: one reference slot per live token, plan modes for the draft depths, +GPU acceptance driving retirement of rejected candidates and their dependent +descendants, same-position replacement through incarnations, delayed +completion and reference-slot reuse. +""" + +import math + +import pytest +import torch + +from tests.helpers import ( + bootstrap_speculation, + make_cpu_cache, + make_speculative_engine, + tiny_ouro_model, +) +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention import sparse as sparse_module +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + REFERENCE_DEPTH, + SparseDecodePlan, + reference_mode, + reference_statistics, + retained_count, + sparse_correction, + sparse_paged_attention, +) +from vllm_rlt.worker.attention_reference import AttentionReferenceStore +from vllm_rlt.worker.metadata import MetadataBuilder + +tiny = tiny_ouro_model +bootstrap = bootstrap_speculation + + +def drain(engine): + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + return outputs + for output in engine.step(): + outputs[output.request_id] = output + pytest.fail("engine failed to drain") + + +def make_store(**overrides): + values = dict(slots=6, layers=2, heads=2, head_dim=4, key_bucket=8, device="cpu") + values.update(overrides) + return AttentionReferenceStore(**values) + + +def fill_payload(store, slots, value=1.0): + for slot in slots: + store.mass[slot] = value + store.output[slot] = value + store.contribution[slot] = value + store.selected[slot] = 1 + store.counts[slot] = 1 + + +# --------------------------------------------------------------------------- +# Plan and arithmetic contracts (RFC #156 sections 3.2 and 6.3) +# --------------------------------------------------------------------------- + + +def test_reference_mode_follows_reference_depth(): + # Loop 2 (depth index 1) captures; later loops consume; loop 1 stays dense. + assert reference_mode(0) == DENSE + assert reference_mode(REFERENCE_DEPTH) == CAPTURE + assert reference_mode(2) == CONSUME + assert reference_mode(3) == CONSUME + + +def test_plan_groups_rows_by_mode(): + plan = SparseDecodePlan( + modes=(DENSE, CAPTURE, CAPTURE, CONSUME), + slots=(-1, 3, 4, 5), + incarnations=(0, 1, 1, 1), + key_retention=0.5, + key_bucket=None, + ) + assert plan.groups() == {DENSE: [0], CAPTURE: [1, 2], CONSUME: [3]} + assert not plan.dense() + assert plan.reference_depth == REFERENCE_DEPTH + + +def test_retained_count_follows_flashloop_ceil_rule(): + # FlashLoop aeaadee select_decode_keys: k = min(N, max(1, ceil(N * r))). + assert retained_count(13, 0.1) == 2 + assert retained_count(10, 0.1) == 1 + assert retained_count(10, 0.5) == 5 + assert retained_count(9, 0.5) == 5 + assert retained_count(1, 0.1) == 1 + assert retained_count(4, 1.0) == 4 + assert retained_count(0, 0.5) == 0 + with pytest.raises(ValueError, match="key_retention"): + retained_count(4, 1.5) + + +def test_reference_statistics_fail_when_capacity_cannot_hold_retained_keys(): + # Capacity never truncates the algorithm: a count above it fails loudly. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.2, 0.3, 0.5])).reshape(3, 1, 1) + values = torch.ones(3, 1, 1) + with pytest.raises(ValueError, match="cannot hold"): + reference_statistics(q, keys, values, key_retention=0.5, key_bucket=1) + # The same query passes with sufficient capacity. + output, _, _, positions = reference_statistics(q, keys, values, key_retention=0.5, key_bucket=2) + assert positions.shape[1] == 2 + assert output.shape == (1, 1) + + +def test_reference_statistics_match_rfc_worked_example(): + # RFC #156 section 3.2: probs [0.1, 0.3, 0.6], values [10, 20, 30]; + # full output 25; retaining the last two keys gives m=0.9, C_ref=24. + # A later loop with local softmax [0.75, 0.25] over values [20, 40] + # corrects to 25 - 24 + 0.9 * (0.75*20 + 0.25*40) = 23.5. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.1, 0.3, 0.6])).reshape(3, 1, 1) + values = torch.tensor([10.0, 20.0, 30.0]).reshape(3, 1, 1) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=0.5) + assert output.item() == pytest.approx(25.0) + assert mass.item() == pytest.approx(0.9) + assert contribution.item() == pytest.approx(24.0) + assert positions.tolist() == [[1, 2]] + + later_keys = torch.log(torch.tensor([3.0, 1.0])).reshape(1, 2, 1) + later_values = torch.tensor([20.0, 40.0]).reshape(1, 2, 1) + corrected = sparse_correction(q, later_keys, later_values, mass, output, contribution) + assert corrected.item() == pytest.approx(23.5) + + +def test_full_retention_matches_dense_attention(): + torch.manual_seed(7) + length, kv_heads, heads, dim = 9, 2, 4, 8 + q = torch.randn(heads, dim) + keys = torch.randn(length, kv_heads, dim) + values = torch.randn(length, kv_heads, dim) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=1.0) + # Full retention keeps every key, so the correction reproduces the dense + # reference output exactly. + dense = reference_statistics(q, keys, values, key_retention=1.0)[0] + groups = heads // kv_heads + corrected = sparse_correction( + q, + keys.permute(1, 0, 2).repeat_interleave(groups, dim=0), + values.permute(1, 0, 2).repeat_interleave(groups, dim=0), + mass, + output, + contribution, + ) + assert positions.shape[1] == length + assert torch.allclose(corrected, dense, atol=1e-5) + assert torch.allclose(output, dense, atol=1e-5) + + +@pytest.mark.parametrize("retention, expected", [(0.5, 23.5), (1.0, 40.0)]) +def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expected): + store = make_store(slots=1, layers=1, heads=1, head_dim=1, key_bucket=3) + lease = store.reserve("r", 0, 1, generation=1) + q = torch.ones(1, 1, 1) + keys = torch.zeros(1, 4, 1, 1) + values = torch.zeros_like(keys) + keys[0, :3, 0, 0] = torch.log(torch.tensor([0.1, 0.3, 0.6])) + values[0, :3, 0, 0] = torch.tensor([10.0, 20.0, 30.0]) + tables, lengths = torch.tensor([[0]], dtype=torch.int32), torch.tensor([3]) + capture = store.build_plan([(lease, 0)], [CAPTURE], key_retention=retention) + # No dense callback: exercise the actual paged Torch oracle even at r=1. + output = sparse_paged_attention(q, keys, values, tables, lengths, capture, layer=0) + assert output.item() == pytest.approx(25.0) + keys[0, :3, 0, 0] = torch.log(torch.tensor([1.0, 3.0, 1.0])) + values[0, :3, 0, 0] = torch.tensor([100.0, 20.0, 40.0]) + consume = store.build_plan([(lease, 0)], [CONSUME], key_retention=retention) + output = sparse_paged_attention(q, keys, values, tables, lengths, consume, layer=0) + assert output.item() == pytest.approx(expected) + assert store.uses.item() == 1 + store.release(lease) + + +# --------------------------------------------------------------------------- +# Store lifecycle (RFC #156 section 5) +# --------------------------------------------------------------------------- + + +def test_reserve_uses_one_slot_per_live_token(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + assert len(lease) == 3 + assert len(set(lease.slots)) == 3 + assert lease.incarnations == (1, 1, 1) + assert store.outstanding == 3 + assert len(store.free) == 3 + + +def test_release_requires_owner_and_retires_metadata(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + fill_payload(store, lease.slots) + store.release(lease) + for slot in lease.slots: + # Readers are gated by counts/mass/valid. The retained index, output and + # contribution buffers stay allocated and are overwritten by the next + # capture's used prefix, so retiring never writes the whole bucket. + assert int(store.counts[slot].sum()) == 0 + assert store.mass[slot].sum() == 0 + assert not bool(store.valid[slot].any()) + assert store.outstanding == 0 + # A retired lease cannot release its slots a second time. + with pytest.raises(RuntimeError, match="different owner"): + store.release(lease) + + +def test_accept_invalidates_rejected_candidates_and_descendants(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + store.accept([lease], torch.tensor([1])) + accepted, middle, rejected = lease.slots + # Rows 0..1 survive the accepted prefix; the rejected candidate and every + # descendant stop being readable. + assert int(store.counts[accepted].max()) > 0 and int(store.counts[middle].max()) > 0 + assert int(store.counts[rejected].max()) == 0 + assert store.mass[rejected].sum() == 0 + assert not bool(store.valid[rejected].any()) + assert int(store.accepted[accepted]) == 1 + rejected_plan = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected_plan, layer=0, row=0) + + +def test_accept_drops_host_counts_so_rejected_rows_cannot_size_a_read(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + for slot in lease.slots: + for layer in range(store.layers): + store.note_count(slot, layer, 1) + store.accept([lease], torch.tensor([1])) + assert all(store.host_count(slot, 0) is None for slot in lease.slots) + # The kept row still reads its device count; the rejected row cannot be read. + kept = store.build_plan([(lease, 1)], [CONSUME], key_retention=0.5) + assert store.read_reference(kept, layer=0, row=0)[0].shape[-1] == 1 + rejected = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected, layer=0, row=0) + + +def test_accept_keeps_root_when_every_candidate_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + store.accept([lease], torch.tensor([0])) + assert store.mass[lease.slots[0]].sum() > 0 + for slot in lease.slots[1:]: + assert store.mass[slot].sum() == 0 + + +def test_capture_and_consume_mark_layer_state(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + plan = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + assert not bool(store.valid[lease.slots[0], 0]) + store.capture(plan, layer=0) + store.capture(plan, layer=1) + assert bool(store.valid[lease.slots[0], 0]) and bool(store.valid[lease.slots[0], 1]) + assert not bool(store.valid[lease.slots[1], 0]) + store.consume(plan, layer=0) + assert int(store.uses[lease.slots[1], 0]) == 1 + assert int(store.uses[lease.slots[0], 0]) == 0 + + +def test_stale_plan_after_reincarnation_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + stale = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + store.release(lease) + # Same request and rows, a fresh token incarnation. + renewed = store.reserve("r", 0, 2, generation=1) + assert all(new > old for new, old in zip(renewed.incarnations, lease.incarnations)) + with pytest.raises(RuntimeError, match="stale"): + store.capture(stale, layer=0) + current = store.build_plan([(renewed, 0), (renewed, 1)], [CAPTURE, CONSUME], key_retention=0.5) + store.capture(current, layer=0) + + +def test_outstanding_rounds_hold_disjoint_slots(): + store = make_store(slots=4) + first = store.reserve("a", 0, 2, generation=1) + second = store.reserve("b", 1, 2, generation=1) + assert not set(first.slots) & set(second.slots) + with pytest.raises(RuntimeError, match="no retired attention reference slots"): + store.reserve("c", 2, 1, generation=1) + store.release(first) + third = store.reserve("c", 2, 1, generation=1) + assert set(third.slots) <= set(first.slots) + + +def test_reincarnation_bumps_on_slot_reuse(): + store = make_store() + lease = store.reserve("r", 0, 1, generation=1) + slot = lease.slots[0] + store.release(lease) + again = store.reserve("r", 0, 1, generation=1) + assert again.slots[0] == slot + assert again.incarnations[0] == lease.incarnations[0] + 1 + assert int(store.incarnation[slot]) == again.incarnations[0] + + +def test_bytes_matches_layout(): + store = make_store(slots=3, layers=2, heads=2, head_dim=4, key_bucket=8) + expected = 3 * 2 * 2 * (8 + 1 + 1 + 4 + 4) * 4 + 3 * 2 * (1 + 4) + 3 * (8 + 8 + 4) + assert store.bytes == expected + + +def test_sparse_config_validates_serving_context_bound(): + assert SparseDecodeAttentionConfig(max_context_tokens=512).max_context_tokens == 512 + assert SparseDecodeAttentionConfig().max_context_tokens is None + with pytest.raises(ValueError, match="max_context_tokens"): + SparseDecodeAttentionConfig(max_context_tokens=0) + + +def test_device_metadata_carries_reference_plan(): + class StubCache: + device = torch.device("cpu") + + def _require_live_batch(self, batch): + return None + + builder = MetadataBuilder(StubCache()) + plan = SparseDecodePlan( + modes=(CAPTURE, CONSUME), + slots=(0, 1), + incarnations=(1, 1), + key_retention=0.5, + key_bucket=4, + ) + common = dict( + allocations=(("r", object()),), + position_ids=torch.zeros(2, dtype=torch.int64), + write_blocks=torch.zeros(2, dtype=torch.int64), + write_offsets=torch.zeros(2, dtype=torch.int64), + block_tables=torch.zeros((2, 1), dtype=torch.int32), + context_lengths=torch.zeros(2, dtype=torch.int32), + ) + batch = builder.prepare_device(**common, reference_plan=plan) + assert batch.reference_plan is plan + padded = {**common, "position_ids": torch.zeros(3, dtype=torch.int64)} + with pytest.raises(ValueError, match="one row mode per query row"): + builder.prepare_device(**padded, reference_plan=plan) + + +# --------------------------------------------------------------------------- +# End-to-end decode through the model, the paged KV cache and attention +# --------------------------------------------------------------------------- + + +def step_metadata(cache, request_id, position, depth, plan=None): + """Device-style recurrent metadata for one token at one loop depth.""" + allocation = cache._get_allocation(request_id) + page = position // cache.block_size + table = torch.tensor([allocation.block_tables[depth][: page + 1]], dtype=torch.int32) + return cache.metadata_builder.prepare_device( + allocations=((request_id, allocation),), + position_ids=torch.tensor([position], dtype=torch.int64), + write_blocks=table[:, page].long(), + write_offsets=torch.tensor([position % cache.block_size], dtype=torch.int64), + block_tables=table, + context_lengths=torch.tensor([position + 1], dtype=torch.int32), + reference_plan=plan, + ) + + +def decode(model, cache, request_id, tokens, *, store=None, lease=None, retention=1.0): + """Serial greedy decode: prelude, every loop depth, coda, next token.""" + assert cache.allocate(request_id, len(tokens) + 1) + logits = [] + for position, token in enumerate(tokens): + hidden = model.prelude(torch.tensor([token])) + for depth in range(model.config.total_ut_steps): + plan = None + mode = reference_mode(depth) + if store is not None and mode != DENSE: + plan = store.build_plan([(lease, 0)], [mode], key_retention=retention) + metadata = step_metadata(cache, request_id, position, depth, plan) + hidden, _ = model.recurrent_prepared(hidden, metadata, cache, compute_gate=False) + logits.append(model.coda(hidden)) + return logits + + +def make_reference_store(model, **overrides): + values = dict( + slots=1, + layers=model.config.num_hidden_layers, + heads=model.config.num_attention_heads, + head_dim=model.config.head_dim, + key_bucket=model.config.max_position_embeddings, + device="cpu", + ) + values.update(overrides) + return AttentionReferenceStore(**values) + + +def test_sparse_decode_matches_dense_through_model_and_kv(): + # Full retention is the RFC's control: the corrected output must reproduce + # the dense path through the real model, paged KV cache and attention. + model = tiny_ouro_model() + tokens = [2, 3, 4, 5, 7] + dense = decode(model, make_cpu_cache(model.config), "dense", tokens) + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + sparse = decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=1.0, + ) + assert store.uses.sum() > 0, "the sparse path never consumed a reference" + for expected, actual in zip(dense, sparse): + assert torch.allclose(actual, expected, atol=1e-4), float((actual - expected).abs().max()) + + +def test_sparse_decode_diverges_as_retention_shrinks(): + model = tiny_ouro_model() + tokens = [2, 3, 4, 5, 7, 9, 11] + dense = decode(model, make_cpu_cache(model.config), "dense", tokens) + errors = {} + for retention in (1.0, 0.5, 0.25): + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + sparse = decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=retention, + ) + errors[retention] = max( + float((actual - expected).abs().max()) for expected, actual in zip(dense, sparse) + ) + assert errors[1.0] < 1e-4 + assert errors[0.25] > errors[1.0] + assert errors[0.5] <= errors[0.25] * 4 + + +def test_sparse_consume_reads_only_retained_keys(monkeypatch): + model = tiny_ouro_model() + tokens = [2, 3, 4, 5] + read_counts = [] + original = sparse_module.gather_retained + + def recording(*args, **kwargs): + keys, values = original(*args, **kwargs) + read_counts.append(keys.shape[1]) + return keys, values + + monkeypatch.setattr(sparse_module, "gather_retained", recording) + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=0.5, + ) + # Two consuming loops per token, each reading only the retained budget, and + # every layer gathers its own retained keys. + per_token = model.config.num_hidden_layers * 2 + expected = [ + retained_count(position + 1, 0.5) + for position in range(len(tokens)) + for _ in range(per_token) + ] + assert read_counts == expected + assert max(read_counts) < len(tokens) + + +# --------------------------------------------------------------------------- +# Async speculative integration (GPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_async_sparse_reference_plans_cover_capture_and_consume(): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + assert store is not None + observed = [] + original = store.build_plan + + def recording(owners, modes, **kwargs): + plan = original(owners, modes, **kwargs) + observed.append(plan.modes) + return plan + + store.build_plan = recording + bootstrap(e) + drain(e) + assert any(CAPTURE in modes for modes in observed) + assert any(CONSUME in modes for modes in observed) + # Dense steps never publish a plan. + assert all(DENSE not in modes for modes in observed) + assert store.outstanding == 0 + assert len(store.free) == store.slots + + +@pytest.mark.gpu +@pytest.mark.parametrize("accepted", [0, 1, 2]) +def test_async_sparse_reference_rejection_invalidates_rejected_rows(accepted, monkeypatch): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + bootstrap(e) + calls = 0 + + def controlled(hidden): + nonlocal calls + phase = calls % 4 + calls += 1 + logits = torch.full((len(hidden), model.config.vocab_size), -100.0, device="cuda") + if phase < 3: + logits[:, 10 + phase] = 100 + else: + for row in range(4): + logits[row, 10 + row if row < accepted else 20 + row] = 100 + return logits + + monkeypatch.setattr(e.speculative_runner, "_coda", controlled) + + observed = [] + original = store.accept + + def recording(leases, accepted_counts): + # Pretend the reference loop published state, then let the device-side + # retirement decide which rows keep it. + for lease in leases: + fill_payload(store, lease.slots) + original(leases, accepted_counts) + torch.cuda.synchronize() + counts = [int(a) for a in accepted_counts.tolist()] + for lease, count in zip(leases, counts): + observed.append( + ( + count, + [float(store.mass[slot, 0, 0]) for slot in lease.slots], + [int(store.accepted[slot]) for slot in lease.slots], + ) + ) + + store.accept = recording + e.step() + e.step() + e._collect_speculative() + assert observed, "acceptance-driven retirement never ran" + for count, masses, limits in observed: + assert count == accepted + assert all(limit == accepted for limit in limits) + for row, mass in enumerate(masses): + assert mass == (1.0 if row <= accepted else 0.0) + e.abort_request("r") + drain(e) + assert store.outstanding == 0 + assert len(store.free) == store.slots + + +@pytest.mark.gpu +def test_async_sparse_reference_incarnations_advance_across_rounds(): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + bootstrap(e) + incarnations = [] + original = store.build_plan + + def recording(owners, modes, **kwargs): + plan = original(owners, modes, **kwargs) + incarnations.append(plan.incarnations) + return plan + + store.build_plan = recording + e.step() + e.step() + e._collect_speculative() + assert len(incarnations) >= 2 + # Every new round is a new token incarnation, even at the same position. + assert incarnations[-1][0] > incarnations[0][0] + e.abort_request("r") + drain(e) + assert store.outstanding == 0 + + +@pytest.mark.gpu +def test_sparse_engine_matches_dense_tokens_at_full_retention(): + # End-to-end through the serving engine: full retention must reproduce the + # dense engine's tokens exactly, including async speculative rounds. + def run(sparse): + model = tiny("cuda") + engine = make_speculative_engine(model, sparse=sparse) + prompts = [[2, 3, 4], [5, 6]] + for i, prompt in enumerate(prompts): + engine.add_request( + str(i), prompt, SamplingParams(max_tokens=12 - 4 * i, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + outputs[output.request_id] = output + tokens = [outputs[str(i)].token_ids for i in range(len(prompts))] + outstanding = engine.reference_store.outstanding if engine.reference_store else 0 + engine.close() + return tokens, outstanding + + dense, _ = run(None) + sparse, outstanding = run(SparseDecodeAttentionConfig(enabled=True, key_retention=1.0)) + assert sparse == dense + assert outstanding == 0 + + +@pytest.mark.gpu +def test_release_from_another_stream_builds_its_index_on_the_owner_stream(monkeypatch): + # A speculative ticket retires from the host's current stream while the + # store clears on its owner stream. An index built on the current stream is + # freed when release returns and can be reused there before the owner + # stream reads it, clearing some other slot. + owner = torch.cuda.Stream() + store = AttentionReferenceStore( + slots=8, layers=1, heads=1, head_dim=4, key_bucket=4, device="cuda", stream=owner + ) + lease = store.reserve("r", 0, 2, generation=1) + other = next(slot for slot in range(store.slots) if slot not in lease.slots) + store.valid.fill_(True) + torch.cuda.synchronize() + clear, streams = store._clear, [] + + def record(indices): + streams.append(torch.cuda.current_stream()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + with torch.cuda.stream(owner): + torch.cuda._sleep(20_000_000) + store.release(lease) + poison = [torch.full((1,), other, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert streams == [owner] and len(poison) == 32 + assert not bool(store.valid[list(lease.slots)].any()) + assert bool(store.valid[other].all()) + + +@pytest.mark.gpu +def test_accept_and_retire_do_not_allocate_bucket_sized_temporaries(): + # Reviewer-reported OOM: acceptance used to read-modify-write the retained + # index buffer, materialising a bucket-sized temporary per batch. Retirement + # must stay bounded by the rejected rows, not by slot count x bucket. + torch.cuda.empty_cache() + store = AttentionReferenceStore( + slots=16, layers=24, heads=16, head_dim=128, key_bucket=32768, device="cuda" + ) + lease = store.reserve("r", 0, 16, generation=1) + store.counts[list(lease.slots)] = 1 + store.mass[list(lease.slots)] = 1.0 + store.output[list(lease.slots)] = 1.0 + store.contribution[list(lease.slots)] = 1.0 + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + store.accept([lease], torch.tensor([0], device="cuda")) + store.release(lease) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() + bucket_bytes = 16 * 24 * 16 * 32768 * 4 + assert bucket_bytes > 512 * 1024**2, "test needs a bucket large enough to detect the bug" + assert peak - baseline < bucket_bytes // 8, ( + f"retirement allocated {(peak - baseline) / 1024**2:.1f} MiB " + f"for a {bucket_bytes / 1024**2:.1f} MiB bucket" + ) + + +@pytest.mark.gpu +def test_sparse_engine_bounds_reference_pool_by_serving_context(): + model = tiny("cuda") + e = make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, max_context_tokens=64), + ) + # The pool follows the supported serving bound, not the model's 128-token + # positional limit. + assert e.reference_store.key_bucket == 64 + e.close() + + +@pytest.mark.gpu +def test_sparse_engine_rejects_context_beyond_model_limit(): + model = tiny("cuda") + with pytest.raises(ValueError, match="exceeds the model context"): + make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, max_context_tokens=4096) + ) + + +@pytest.mark.gpu +def test_sparse_async_overlapping_rounds_keep_isolated_references(): + # Reviewer-requested coverage: two submitted rounds of the same request are + # in flight at once. They must hold disjoint reference slots, and retiring + # the older round's rejected candidates must not disturb the newer round. + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5) + ) + store = e.reference_store + bootstrap(e) + accepted_rounds = [] + original = store.accept + + def recording(leases, accepted): + accepted_rounds.append(tuple(lease.slots for lease in leases)) + return original(leases, accepted) + + store.accept = recording + e.step() + e.step() + assert len(e._pending_speculative) == 2, "expected two overlapping rounds" + assert len(accepted_rounds) >= 2 + first, second = accepted_rounds[0][0], accepted_rounds[1][0] + assert first and second + assert not set(first) & set(second), "outstanding rounds share reference slots" + assert store.outstanding == len(first) + len(second) + e._collect_speculative() + # Only the older round retired; the newer round's references survive. + assert store.outstanding == len(second) + e.abort_request("r") + drain(e) + assert store.outstanding == 0 + assert len(store.free) == store.slots + e.close() + + +@pytest.mark.gpu +def test_sparse_engine_sizes_key_buffer_from_retention(): + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.25) + ) + # None key_bucket derives the capacity from context x retention, not the + # full model context (FlashLoop ceil rule at the worst-case key count). + assert e.reference_store.key_bucket == math.ceil(model.config.max_position_embeddings * 0.25) + + +@pytest.mark.gpu +def test_sparse_engine_rejects_insufficient_key_bucket(): + model = tiny("cuda") + with pytest.raises(ValueError, match="key_bucket"): + make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, key_bucket=4), + ) + + +@pytest.mark.gpu +def test_sparse_engine_e2e_multiple_requests_abort_and_drain(): + # Engine-level lifecycle e2e: concurrent rounds across requests, an abort + # in flight, and full slot reclamation without stale references. + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5) + ) + store = e.reference_store + max_outstanding = 0 + original_reserve = store.reserve + + def counting_reserve(*args, **kwargs): + nonlocal max_outstanding + lease = original_reserve(*args, **kwargs) + max_outstanding = max(max_outstanding, store.outstanding) + return lease + + store.reserve = counting_reserve + params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in (5, 9, 50)] + for i in range(3): + e.add_request(str(i), [2, 3, 4], params[i]) + outputs = {} + for _ in range(1000): + if outputs.get("0") is not None and outputs["0"].finished: + break + for output in e.step(): + outputs[output.request_id] = output + # Abort one request mid-flight; the rest must finish cleanly. + if "2" not in outputs or not outputs["2"].finished: + e.abort_request("2") + outputs.update(drain(e)) + assert outputs["0"].finished and len(outputs["0"].token_ids) == 5 + assert outputs["1"].finished and len(outputs["1"].token_ids) == 9 + assert store.outstanding == 0 + assert len(store.free) == store.slots + assert max_outstanding > 0 diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py new file mode 100644 index 00000000..17e8bc67 --- /dev/null +++ b/tests/test_sparse_decode_batched.py @@ -0,0 +1,167 @@ +"""Batched CUDA arithmetic versus the original row-wise sparse oracle.""" + +import pytest +import torch + +from tests.test_sparse_decode_attention import make_store +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + batched_paged_attention, + reference_paged_attention, +) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("use_dense_callback", [False, True]) +def test_batched_ragged_gqa_matches_scalar_oracle( + device, dtype, retention, ties, use_dense_callback +): + torch.manual_seed(157) + rows, heads, kv_heads, dim, block_size = 6, 4, 2, 128, 16 + lengths = torch.tensor([0, 17, 17, 31, 31, 257], device=device) + tables = torch.randperm(rows * 17, device=device).reshape(rows, 17).int() + keys = torch.randn(rows * 17, block_size, kv_heads, dim, device=device, dtype=dtype) + values = torch.randn_like(keys) + queries = torch.randn(rows, heads, dim, device=device, dtype=dtype) + if ties: + queries.zero_() + stores = [ + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=257, + device=device, + initial_key_capacity=1 if implementation else None, + ) + for implementation in range(2) + ] + leases = [store.reserve("test", 0, rows, generation=1) for store in stores] + implementations = [reference_paged_attention, batched_paged_attention] + for modes in ([CAPTURE] * rows, [DENSE, CAPTURE, CONSUME, CONSUME, DENSE, CONSUME]): + results = [] + for store, lease, implementation in zip(stores, leases, implementations): + plan = store.build_plan( + [(lease, row) for row in range(rows)], modes, key_retention=retention + ) + # A deliberately distinct backend output verifies that capture + # stores the rounded recurrence input, not its FP32 oracle PV. + callback = (lambda q, *args: q * 0.375 + 1.25) if use_dense_callback else None + results.append( + implementation( + queries, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=callback, + ) + ) + torch.testing.assert_close(results[1], results[0], rtol=0, atol=0) + capacity = stores[1].selected.shape[-1] + assert torch.equal(stores[0].selected[..., :capacity], stores[1].selected) + assert not torch.count_nonzero(stores[0].selected[..., capacity:]) + for name in ("counts", "valid", "uses"): + assert torch.equal(getattr(stores[0], name), getattr(stores[1], name)), name + for name in ("mass", "output", "contribution"): + torch.testing.assert_close( + getattr(stores[0], name), getattr(stores[1], name), rtol=0, atol=0 + ) + queries = queries + 0.25 + keys = keys + 0.125 + values = values - 0.125 + for store, lease in zip(stores, leases): + store.release(lease) + assert store.outstanding == 0 + + +def test_batched_rejects_missing_and_stale_reference(): + store = make_store(slots=1, layers=1, heads=2, head_dim=4) + lease = store.reserve("test", 0, 1, generation=1) + plan = store.build_plan([(lease, 0)], [CONSUME], key_retention=0.1) + args = ( + torch.ones(1, 2, 4), + torch.ones(1, 16, 2, 4), + torch.ones(1, 16, 2, 4), + torch.zeros(1, 1, dtype=torch.int32), + torch.tensor([1]), + plan, + ) + with pytest.raises(RuntimeError, match="no captured reference"): + batched_paged_attention(*args, layer=0) + store.release(lease) + store.reserve("test", 0, 1, generation=2) + with pytest.raises(RuntimeError, match="stale attention reference"): + batched_paged_attention(*args, layer=0) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_reference_growth_preserves_live_layers_and_budget(device): + store = make_store( + slots=3, + layers=2, + heads=2, + head_dim=4, + key_bucket=17, + initial_key_capacity=2, + device=device, + ) + leases = [store.reserve(str(i), i, 1, generation=1) for i in range(2)] + plans = [store.build_plan([(lease, 0)], [CAPTURE], key_retention=1.0) for lease in leases] + # Simulate queued work on the runner's owner stream before reallocating. + if device == "cuda": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + store.stream = stream + context = torch.cuda.stream(stream) + else: + from contextlib import nullcontext + + context = nullcontext() + budget = store.peak_capacity_bytes + fixed = store.bytes - store.selected.nbytes + assert budget == fixed + 3 * 2 * 2 * (2 + 4 + 8 + 16 + 17) * 4 + with context: + for layer in range(2): + store.write_reference( + plans[0], + layer, + 0, + selected=torch.tensor([[1, 3], [2, 4]], device=device), + counts=torch.full((2,), 2, device=device), + mass=torch.full((2,), 0.5, device=device), + output=torch.full((2, 4), 0.25, device=device), + contribution=torch.full((2, 4), 0.125, device=device), + ) + before = [x.clone() for x in store.read_reference(plans[0], 1, 0)] + # Grow several times while a different lease remains live. + for count in [3, 7, 17]: + store.ensure_key_capacity(count) + assert store.selected.shape[-1] >= count + assert store.peak_capacity_bytes == budget + for actual, expected in zip(store.read_reference(plans[0], 1, 0), before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + with pytest.raises(ValueError, match="capacity"): + store.ensure_key_capacity(18) + if device == "cuda": + store.stream.synchronize() + for lease in leases: + store.release(lease) + assert store.outstanding == 0 + + +@pytest.mark.gpu +def test_reference_growth_rejects_foreign_stream(): + store = make_store(key_bucket=17, initial_key_capacity=2, device="cuda") + store.stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="owner stream"): + store.ensure_key_capacity(3) + assert store.selected.shape[-1] == 2 diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py new file mode 100644 index 00000000..ecae5374 --- /dev/null +++ b/tests/test_sparse_decode_fused.py @@ -0,0 +1,633 @@ +"""Opt-in fused sparse decode kernels versus the batched oracle-shaped path.""" + +import pytest +import torch + +from tests.helpers import make_speculative_engine, tiny_ouro_model +from tests.test_sparse_decode_attention import drain, make_store +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + reference_paged_attention, + sparse_paged_attention, +) + + +def test_sparse_config_validates_implementation(): + assert SparseDecodeAttentionConfig().implementation == "batched" + assert SparseDecodeAttentionConfig(implementation="fused").implementation == "fused" + with pytest.raises(ValueError, match="implementation"): + SparseDecodeAttentionConfig(implementation="native") + + +def test_fused_store_keeps_the_oracle_on_cpu(monkeypatch): + calls = [] + import vllm_rlt.attention.sparse as sparse + + monkeypatch.setattr(sparse, "reference_paged_attention", lambda *a, **k: calls.append(1)) + store = make_store(slots=1, layers=1, heads=2, head_dim=4, implementation="fused") + lease = store.reserve("cpu", 0, 1, generation=1) + plan = store.build_plan([(lease, 0)], [CAPTURE], key_retention=0.5) + q = torch.ones(1, 2, 4) + cache = torch.ones(1, 16, 2, 4) + sparse.sparse_paged_attention( + q, cache, cache, torch.zeros(1, 1, dtype=torch.int32), torch.tensor([3]), plan, layer=0 + ) + assert calls == [1] + + +@pytest.fixture(autouse=True) +def _check_fast_launches(monkeypatch): + """Every reused launch must be the kernel Triton's own dispatch would pick.""" + try: + from vllm_rlt.kernels import sparse_decode + except ImportError: # CPU-only installs have no Triton and no fused launches + return + monkeypatch.setattr(sparse_decode, "_LAUNCH_CHECK", True) + + +@pytest.mark.gpu +@pytest.mark.parametrize("context", [37, 1024, 4111]) +def test_fast_launches_match_triton_dispatch_bit_for_bit(context, monkeypatch): + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + if not sparse_decode._FAST_LAUNCH: + pytest.skip("fast launches are limited to Triton 3.8") + rows, heads, kv_heads, dim, block_size, layers = 3, 4, 2, 64, 16, 18 + pages = -(-context // block_size) + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.bfloat16, 3 + ) + lengths = torch.tensor([context, max(1, context - 17), max(1, context // 2)], device="cuda") + results = {} + for fast in (False, True): + monkeypatch.setattr(sparse_decode, "_FAST_LAUNCH", fast) + store = make_store( + slots=rows, + layers=layers, + heads=heads, + head_dim=dim, + key_bucket=context, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + owners = [(lease, row) for row in range(rows)] + outputs = [] + # Layers 0, 1, 16 and 17 differ in Triton's integer specialization. + for _ in range(2): + for layer in (0, 1, 2, 16, 17): + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.1) + outputs.append( + sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=layer, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + ) + torch.cuda.synchronize() + results[fast] = ( + outputs, + store.selected.clone(), + store.counts.clone(), + store.mass.clone(), + store.contribution.clone(), + store.output.clone(), + ) + slow, fast = results[False], results[True] + for a, b in zip(slow[0], fast[0]): + assert torch.equal(a, b) + for a, b in zip(slow[1:], fast[1:]): + assert torch.equal(a, b) + + +@pytest.mark.gpu +def test_consume_scratch_follows_the_head_dimension(monkeypatch): + """A stream's consume scratch must not be reused across head dimensions.""" + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + def consume(dim): + rows, heads, block_size, pages = 3, 4, 16, 5 + q, keys, values, tables = _paged(rows, heads, 2, dim, pages, block_size, torch.float32, 4) + lengths = torch.tensor([7, 33, 80], device="cuda", dtype=torch.int32) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=80, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + owners = [(lease, row) for row in range(rows)] + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.25) + output = sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + torch.cuda.synchronize() + return output + + monkeypatch.setattr(sparse_decode, "_CONSUME_WORKSPACE", {}) + fresh = consume(128) + sparse_decode._CONSUME_WORKSPACE.clear() + consume(16) + # The same stream now holds 16-wide scratch; a 128-wide consume must regrow it. + assert torch.equal(consume(128), fresh) + + +def _paged(rows, heads, kv_heads, dim, pages, block_size, dtype, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + tables = torch.randperm(rows * pages, device="cuda", generator=generator) + keys = torch.randn( + rows * pages, block_size, kv_heads, dim, device="cuda", dtype=dtype, generator=generator + ) + values = torch.randn(keys.shape, device="cuda", dtype=dtype, generator=generator) + queries = torch.randn(rows, heads, dim, device="cuda", dtype=dtype, generator=generator) + return queries, keys, values, tables.reshape(rows, pages).int() + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.5]) +@pytest.mark.parametrize("ties", [False, True]) +# The Triton backend takes the split capture; any other callable takes the fallback. +@pytest.mark.parametrize("split", [True, False]) +def test_fused_capture_and_consume_track_batched_path(retention, ties, split): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention + + backend = TritonPagedAttention() if split else (lambda *args: triton_paged_attention(*args)) + + rows, heads, kv_heads, dim, block_size, pages = 6, 4, 2, 128, 16, 17 + # 10 and 30 sit on float rounding edges of ceil(length * 0.1) (30 * 0.1 > 3 in float64). + lengths = torch.tensor([0, 10, 17, 30, 200, 257], device="cuda", dtype=torch.int32) + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=272, + device="cuda", + initial_key_capacity=1, + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [store.reserve("fused", 0, rows, generation=1) for store in stores] + owners = [[(lease, row) for row in range(rows)] for lease in leases] + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.float32, 1 + ) + if ties: + # Every score equal: both paths must keep the smallest positions. + q.zero_() + captured = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CAPTURE] * rows, key_retention=retention) + captured.append( + sparse_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + batched, fused = stores + # Capture rows return (and publish) the dense backend output unchanged. These + # inputs are FP32, where some architectures (sm_90) round the split copy one + # ulp from the dense kernel; the BF16 test below stays bit for bit. + exact = 0 if not split else 1e-6 + torch.testing.assert_close(captured[1], captured[0], rtol=0, atol=exact) + torch.testing.assert_close(fused.output[:, 1], batched.output[:, 1], rtol=0, atol=exact) + torch.testing.assert_close(fused.counts[:, 1], batched.counts[:, 1], rtol=0, atol=0) + # The finalize kernel marks the same rows source-ready as store.capture. + assert torch.equal(fused.valid, batched.valid) + for row, length in enumerate(lengths.tolist()): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + count = int(batched.counts[slot_b, 1, 0]) + for head in range(heads): + # Both paths write ascending positions; ties keep the smaller ones. + expected = batched.selected[slot_b, 1, head, :count].tolist() + assert fused.selected[slot_f, 1, head, :count].tolist() == expected + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + # A later loop reads its own plane: fresh keys and queries, same pages. + q2, keys2, values2, _ = _paged(rows, heads, kv_heads, dim, pages, block_size, torch.float32, 2) + consumed = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CONSUME] * rows, key_retention=retention) + consumed.append( + sparse_paged_attention( + q2, + keys2, + values2, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + torch.testing.assert_close(consumed[1], consumed[0], rtol=1e-4, atol=1e-5) + assert not consumed[1][0].any(), "an empty row stays zero" + assert int(fused.uses[list(leases[1].slots), 1].min()) == 1 + for store, lease in zip(stores, leases): + store.release(lease) + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_radix_selection_keeps_oracle_ties_on_repeated_scores(retention): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # Keys cycle through three vectors, so every head sees three distinct scores + # repeated across the context: the retained set is decided almost entirely + # by the tie rule (smaller positions first), with negative scores included. + # Small integers keep every dot product exact. With random reals the dense + # Triton kernel can score one key a ULP apart at different tile positions + # (seen on sm_89), so the "ties" would not be ties for either path. + rows, heads, dim, block_size, pages = 3, 4, 64, 16, 40 + generator = torch.Generator(device="cuda").manual_seed(11) + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + q = integers(rows, heads, dim) + lengths = torch.tensor([7, 300, 640], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=640, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + for store in stores: + lease = store.reserve("ties", 0, rows, generation=1) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + batched, fused = stores + assert torch.equal(fused.counts, batched.counts) + for slot in range(rows): + for head in range(heads): + count = int(batched.counts[slot, 0, head]) + assert ( + fused.selected[slot, 0, head, :count].tolist() + == batched.selected[slot, 0, head, :count].tolist() + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("folded", [True, False]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_parallel_select_matches_batched_across_many_parts( + ties, retention, folded, monkeypatch +): + import vllm_rlt.kernels.sparse_decode as fused_module + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # The first radix pass runs inside the split kernel or as its own pass. + monkeypatch.setattr(fused_module, "_PASS0_IN_SPLIT", folded) + + # Long rows split into up to 16 parts; with repeated integer scores the + # retained ties span part boundaries, so offsets and tie quotas must carry. + rows, heads, dim, block_size, pages = 4, 4, 64, 16, 320 + generator = torch.Generator(device="cuda").manual_seed(29) + if ties: + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + q = integers(rows, heads, dim) + else: + # Mostly distinct but still exact scores: a one-ULP difference between the + # Triton and batched scores at this length could cross the boundary. + keys = torch.randint( + -8, 9, (rows * pages, block_size, heads, dim), device="cuda", generator=generator + ).float() + q = torch.randint(-8, 9, (rows, heads, dim), device="cuda", generator=generator).float() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + lengths = torch.tensor([1, 129, 2049, 5000], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=5120, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [] + for store in stores: + lease = store.reserve("parts", 0, rows, generation=1) + leases.append(lease) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=1, dense_attention=backend + ) + batched, fused = stores + for row in range(rows): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + assert torch.equal(fused.counts[slot_f, 1], batched.counts[slot_b, 1]) + for head in range(heads): + count = int(batched.counts[slot_b, 1, head]) + assert ( + fused.selected[slot_f, 1, head, :count].tolist() + == batched.selected[slot_b, 1, head, :count].tolist() + ) + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_split_capture_is_the_dense_backend_output_bit_for_bit(dtype): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + rows, heads, kv_heads, dim = 5, 4, 2, 128 + lengths = torch.tensor([0, 1, 33, 515, 1200], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 76, 16, dtype, 5) + backend = TritonPagedAttention() + expected = backend(q, keys, values, tables, lengths) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=1216, + device="cuda", + implementation="fused", + ) + lease = store.reserve("split", 0, rows, generation=1) + plan = store.build_plan([(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=0.1) + out = sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + # The split copy and the folded merge must not drift from the dense kernels: + # bit for bit in BF16; in FP32 some architectures (sm_90) contract the split + # copy's arithmetic differently and round it one ulp apart. + atol = 0 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(out, expected, rtol=0, atol=atol) + # O_ref is the output actually fed into the recurrence: rounded, then FP32. + live = list(lease.slots)[1:] + torch.testing.assert_close(store.output[live, 0], expected[1:].float(), rtol=0, atol=atol) + assert bool(store.valid[list(lease.slots), 0].all()) + + +@pytest.mark.gpu +def test_fused_bf16_consume_matches_scalar_oracle_within_rounding(): + from vllm_rlt.kernels.paged_attention import triton_paged_attention + + rows, heads, kv_heads, dim = 4, 4, 4, 64 + lengths = torch.tensor([5, 64, 129, 300], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 19, 16, torch.bfloat16, 3) + results = [] + for implementation, function in ( + ("fused", sparse_paged_attention), + ("batched", reference_paged_attention), + ): + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=304, + device="cuda", + implementation=implementation, + ) + lease = store.reserve("bf16", 0, rows, generation=1) + owner = [(lease, row) for row in range(rows)] + for modes in ([CAPTURE] * rows, [CONSUME] * rows): + plan = store.build_plan(owner, modes, key_retention=0.1) + out = function( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=triton_paged_attention, + ) + results.append(out.float()) + # Different FP32 reduction order may move a BF16 rounding step. + torch.testing.assert_close(results[0], results[1], rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +def test_fused_consume_without_capture_is_nan_not_silent(): + rows, heads, dim = 2, 2, 64 + store = make_store( + slots=rows, layers=1, heads=heads, head_dim=dim, device="cuda", implementation="fused" + ) + lease = store.reserve("missing", 0, rows, generation=1) + plan = store.build_plan([(lease, 0), (lease, 1)], [CONSUME, CONSUME], key_retention=0.5) + q, keys, values, tables = _paged(rows, heads, heads, dim, 2, 16, torch.float32, 4) + lengths = torch.tensor([0, 9], device="cuda", dtype=torch.int32) + out = sparse_paged_attention(q, keys, values, tables, lengths, plan, layer=0) + assert not out[0].any() and torch.isnan(out[1]).all() + + +@pytest.mark.gpu +def test_fused_engine_rounds_abort_and_drain(monkeypatch): + import vllm_rlt.kernels.sparse_decode as fused_module + + calls = {CAPTURE: 0, CONSUME: 0} + original = fused_module.fused_paged_attention + + def counting(*args, **kwargs): + calls[args[5].modes[0]] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(fused_module, "fused_paged_attention", counting) + # The engine's Triton backend adapter must reach the split capture path. + monkeypatch.setattr(fused_module, "_capture_kernel", None) + model = tiny_ouro_model("cuda") + engine = make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused"), + ) + store = engine.reference_store + assert store.implementation == "fused" + params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in (5, 9, 50)] + for i in range(3): + engine.add_request(str(i), [2, 3, 4], params[i]) + outputs = {} + for _ in range(1000): + if "0" in outputs and outputs["0"].finished: + break + for output in engine.step(): + outputs[output.request_id] = output + if "2" not in outputs or not outputs["2"].finished: + engine.abort_request("2") + outputs.update(drain(engine)) + assert outputs["0"].finished and len(outputs["0"].token_ids) == 5 + assert outputs["1"].finished and len(outputs["1"].token_ids) == 9 + assert calls[CAPTURE] > 0 and calls[CONSUME] > 0 + assert store.outstanding == 0 and len(store.free) == store.slots + assert engine.cache_manager.num_used_blocks == 0 + engine.close() + + +def test_padding_slot_is_allocated_but_never_bound(): + store = make_store(slots=3, padding_slot=True) + assert store.padding_slot == 3 and store.slots == 3 + assert store.valid.shape[0] == store.uses.shape[0] == 4 + lease = store.reserve("r", 0, 3, generation=1) + assert store.padding_slot not in lease.slots + with pytest.raises(RuntimeError, match="no retired"): + store.reserve("s", 1, 1, generation=1) + store.release(lease) + store.discard() + assert sorted(store.free) == [0, 1, 2] + assert make_store(slots=3).padding_slot is None + + +def _graph_engine(*, graphs, max_rows=32): + from tests.helpers import tiny_ouro_config + from vllm_rlt import CacheConfig, ExecutionConfig, SchedulerConfig, SpeculativeConfig + from vllm_rlt.engine.stage_engine import StageEngine + from vllm_rlt.models import OuroForCausalLM + + torch.manual_seed(123) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16) + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(128, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=3, max_num_batched_tokens=16, prefill_chunk_size=8 + ), + execution_config=ExecutionConfig( + async_scheduling=True, + static_buffers=True, + pad_to_power_of_two=True, + cuda_graphs=graphs, + cuda_graph_max_batch_size=max_rows, + ), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=True, key_retention=0.5, implementation="fused" + ), + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("max_rows", [32, 2]) +def test_fused_graphs_replay_sparse_and_match_eager(max_rows, monkeypatch): + import vllm_rlt.kernels.sparse_decode as fused_module + + # Graph and eager captures must both take the split path, not the serial fallback. + monkeypatch.setattr(fused_module, "_capture_kernel", None) + results = [] + for graphs in (False, True): + engine = _graph_engine(graphs=graphs, max_rows=max_rows) + store = engine.reference_store + try: + used = 0 + for i, (prompt_length, output_length) in enumerate([(15, 23), (17, 9), (31, 1)]): + engine.add_request( + str(i), + [2 + i, 3, 4] * (prompt_length // 3) + [5] * (prompt_length % 3), + SamplingParams(max_tokens=output_length, ignore_eos=True), + ) + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + outputs[output.request_id] = output + used = max(used, int(store.uses[: store.slots].sum())) + assert len(outputs) == 3 + results.append({rid: o.token_ids for rid, o in outputs.items()}) + # Consumers ran on live slots, not only on the padding slot. + assert used > 0 + assert store.outstanding == 0 and len(store.free) == store.slots + assert engine.cache_manager.num_used_blocks == 0 + graphs_owner = engine.speculative_runner.graphs + if graphs: + assert graphs_owner.sparse_replays > 0 + if max_rows == 32: + assert graphs_owner.fallbacks == 0 + else: + assert graphs_owner.fallbacks > 0 + finally: + engine.close() + # The same per-row kernels run eager and replayed. + assert results[0] == results[1] + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "sparse", + [ + SparseDecodeAttentionConfig(enabled=True, key_retention=0.5), + SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, implementation="fused"), + ], +) +def test_graphs_reject_sparse_paths_that_cannot_be_replayed(sparse): + from vllm_rlt import ExecutionConfig, SpeculativeConfig + from vllm_rlt.engine.stage_engine import StageEngine + + with pytest.raises(ValueError, match="CUDA graphs requires"): + StageEngine( + tiny_ouro_model("cuda"), + attention_backend="triton", + execution_config=ExecutionConfig(async_scheduling=True, cuda_graphs=True), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ) diff --git a/tests/test_sparse_decode_pipeline.py b/tests/test_sparse_decode_pipeline.py new file mode 100644 index 00000000..9d459e86 --- /dev/null +++ b/tests/test_sparse_decode_pipeline.py @@ -0,0 +1,286 @@ +"""Sparse decode attention with token pipelines: device-decided row modes and slots.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from tests.test_sparse_decode_attention import drain, make_store +from tests.test_sparse_decode_fused import _paged +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE + +PROMPTS = [[2, 3, 4, 5, 6, 7, 8], [9, 10, 11], [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]] +LENGTHS = (23, 9, 17) + + +@pytest.fixture(autouse=True) +def _check_fast_launches(monkeypatch): + """Every reused launch must be the kernel Triton's own dispatch would pick.""" + try: + from vllm_rlt.kernels import sparse_decode + except ImportError: # CPU-only installs have no Triton and no fused launches + return + monkeypatch.setattr(sparse_decode, "_LAUNCH_CHECK", True) + + +@pytest.mark.gpu +def test_mixed_modes_match_uniform_fused_kernels(): + """One launch set with every mode equals the uniform launches row for row.""" + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + rows, heads, kv_heads, dim, block_size, pages = 6, 4, 2, 64, 16, 9 + dense = TritonPagedAttention() + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.float32, 7 + ) + q2, _, _, _ = _paged(rows, heads, kv_heads, dim, pages, block_size, torch.float32, 8) + lengths = torch.tensor([5, 33, 70, 129, 140, 1], device="cuda", dtype=torch.int32) + modes = [DENSE, CAPTURE, CONSUME, CAPTURE, CONSUME, DENSE] + stores, leases = [], [] + for _ in range(2): + store = make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=pages * block_size, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + # The consume rows need a reference from an earlier loop. + earlier = [row for row, mode in enumerate(modes) if mode == CONSUME] + plan = store.build_plan( + [(lease, row) for row in earlier], [CAPTURE] * len(earlier), key_retention=0.25 + ) + sparse_decode.fused_paged_attention( + q[earlier], + keys, + values, + tables[earlier], + lengths[earlier], + plan, + layer=1, + dense_attention=dense, + ) + stores.append(store) + leases.append(lease) + mixed_store, uniform_store = stores + plan = mixed_store.build_device_plan( + [leases[0]], + leases[0].slot_ids, + torch.tensor(modes, device="cuda", dtype=torch.int8), + key_retention=0.25, + ) + mixed = sparse_decode.fused_paged_attention( + q2, keys, values, tables, lengths, plan, layer=1, dense_attention=dense + ) + expected = torch.empty_like(q2) + for mode in (DENSE, CAPTURE, CONSUME): + subset = [row for row, m in enumerate(modes) if m == mode] + if mode == DENSE: + expected[subset] = dense(q2[subset], keys, values, tables[subset], lengths[subset]) + continue + host = uniform_store.build_plan( + [(leases[1], row) for row in subset], [mode] * len(subset), key_retention=0.25 + ) + expected[subset] = sparse_decode.fused_paged_attention( + q2[subset], + keys, + values, + tables[subset], + lengths[subset], + host, + layer=1, + dense_attention=dense if mode == CAPTURE else None, + ) + torch.cuda.synchronize() + # Consume rows are the uniform consume's, bit for bit. Dense and capture + # rows are the dense kernel's up to fp32 rounding: the mixed-mode variant is + # compiled separately, and some architectures (sm_89) round it differently + # by an ulp. Every store field matches slot for slot. + consume = [row for row, mode in enumerate(modes) if mode == CONSUME] + assert torch.equal(mixed[consume], expected[consume]) + torch.testing.assert_close(mixed, expected, rtol=0, atol=1e-6) + for name in ("selected", "counts", "mass", "output", "contribution", "valid", "uses"): + assert torch.equal(getattr(mixed_store, name), getattr(uniform_store, name)), name + + +@pytest.mark.gpu +def test_steps_without_capture_skip_selection_kernels(monkeypatch): + """A device plan that cannot capture launches four kernels, with identical results.""" + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + rows, heads, kv_heads, dim, block_size, pages = 4, 4, 2, 64, 16, 9 + dense = TritonPagedAttention() + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.float32, 3 + ) + lengths = torch.tensor([40, 129, 7, 90], device="cuda", dtype=torch.int32) + modes = torch.tensor([DENSE, CONSUME, DENSE, CONSUME], device="cuda", dtype=torch.int8) + results, launches = [], [] + for may_capture in (True, False): + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=pages * block_size, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + consume_rows = [1, 3] + capture = store.build_plan( + [(lease, row) for row in consume_rows], [CAPTURE] * 2, key_retention=0.25 + ) + sparse_decode.fused_paged_attention( + q[consume_rows], + keys, + values, + tables[consume_rows], + lengths[consume_rows], + capture, + layer=0, + dense_attention=dense, + ) + plan = store.build_device_plan( + [lease], lease.slot_ids, modes, key_retention=0.25, may_capture=may_capture + ) + count = [] + launch = sparse_decode._launch + monkeypatch.setattr( + sparse_decode, + "_launch", + lambda kernel, *a, **k: (count.append(kernel), launch(kernel, *a, **k)), + ) + output = sparse_decode.fused_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=dense + ) + monkeypatch.setattr(sparse_decode, "_launch", launch) + torch.cuda.synchronize() + launches.append(len(count)) + results.append((output, store.uses.clone(), store.counts.clone(), store.valid.clone())) + assert launches == [8, 4] + for a, b in zip(*results): + assert torch.equal(a, b) + + +def _engine(mode, sparse, *, draft_loops=2, graphs=False, retention=0.5): + from vllm_rlt.engine.stage_engine import StageEngine + from vllm_rlt.models import OuroForCausalLM + + torch.manual_seed(123) + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.float32) + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(256, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=3, max_num_batched_tokens=16, prefill_chunk_size=8 + ), + execution_config=ExecutionConfig( + async_scheduling=True, + static_buffers=graphs, + pad_to_power_of_two=graphs, + cuda_graphs=graphs, + ), + speculative_config=SpeculativeConfig(3, draft_loops=draft_loops, mode=mode), + sparse_decode_attention_config=( + SparseDecodeAttentionConfig( + enabled=True, key_retention=retention, implementation="fused" + ) + if sparse + else None + ), + ) + + +def _generate(e): + for i, (prompt, length) in enumerate(zip(PROMPTS, LENGTHS)): + e.add_request(str(i), prompt, SamplingParams(max_tokens=length, ignore_eos=True)) + outputs = drain(e) + assert sorted(outputs) == ["0", "1", "2"] + assert e.cache_manager.num_used_blocks == 0 + if e.reference_store is not None: + assert e.reference_store.outstanding == 0 + return {rid: list(output.token_ids) for rid, output in outputs.items()} + + +@pytest.mark.gpu +@pytest.mark.parametrize("draft_loops", [1, 2]) +@pytest.mark.parametrize("graphs", [False, True]) +def test_pipeline_sparse_matches_block_sparse(draft_loops, graphs): + """Both capture at loop 1 and consume at loops 2 and 3 of every token.""" + results = {} + for mode in ("block", "pipeline"): + for sparse in (False, True): + e = _engine(mode, sparse, draft_loops=draft_loops, graphs=graphs) + try: + results[mode, sparse] = _generate(e) + if sparse and mode == "pipeline" and graphs: + runner = e.speculative_runner + assert runner.graphs.sparse_replays > 0 + finally: + e.close() + # Precondition: the dense pipeline already reproduces block speculation. + assert results["pipeline", False] == results["block", False] + assert results["pipeline", True] == results["block", True] + + +@pytest.mark.gpu +def test_pipeline_sparse_slots_follow_requests_through_abort_and_reuse(): + e = _engine("pipeline", True) + store = e.reference_store + width = min(4, 4 // 2) + try: + e.add_request("a", PROMPTS[0], SamplingParams(max_tokens=40, ignore_eos=True)) + for _ in range(200): + e.step() + state = e.speculative_runner.states.get("a") + if state is not None and state.reference is not None: + break + else: + pytest.fail("pipeline state was not created") + # One lease of `width` slots for the request's lifetime. + assert store.outstanding == width + assert sorted(state.reference.slots) == sorted(state.pipeline.slots.tolist()) + e.abort_request("a") + drain(e) + assert store.outstanding == 0 + # A reused id gets a new generation and fresh slots. + assert _generate(e) + finally: + e.close() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "sparse", + [ + SparseDecodeAttentionConfig(enabled=True, key_retention=0.5), + SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, implementation="fused"), + ], +) +def test_pipeline_rejects_sparse_paths_without_device_plans(sparse): + from vllm_rlt.engine.stage_engine import StageEngine + from vllm_rlt.models import OuroForCausalLM + + model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.float32) + with pytest.raises(ValueError, match="token pipelines"): + StageEngine( + model, + attention_backend="triton", + execution_config=ExecutionConfig(async_scheduling=True), + speculative_config=SpeculativeConfig(3, mode="pipeline"), + sparse_decode_attention_config=sparse, + ) diff --git a/tests/test_sparse_decode_preemption.py b/tests/test_sparse_decode_preemption.py new file mode 100644 index 00000000..a8ef0d1f --- /dev/null +++ b/tests/test_sparse_decode_preemption.py @@ -0,0 +1,267 @@ +"""Preemption of ordinary sparse decode: a suspended token keeps its reference.""" + +import pytest +import torch + +from tests.test_sparse_decode_sync import _bounded_pool, _ordinary, engine, generate +from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig, SpeculativeConfig +from vllm_rlt.attention.sparse import CAPTURE, CONSUME + + +def _scheduler(**options): + values = dict( + max_num_seqs=1, max_num_batched_tokens=16, prefill_chunk_size=3, enable_preemption=True + ) + values.update(options) + return SchedulerConfig(**values) + + +def _record_modes(e, monkeypatch): + """Record (request, mode) of every sparse plan the runner builds.""" + seen = [] + build = e.reference_store.build_plan + + def spy(owners, modes, **kwargs): + seen.append([(None if o is None else o[0].request_id, m) for o, m in zip(owners, modes)]) + return build(owners, modes, **kwargs) + + monkeypatch.setattr(e.reference_store, "build_plan", spy) + return seen + + +def _drain(e, outputs): + for _ in range(1000): + if not e.has_unfinished_requests(): + return outputs + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + raise AssertionError("engine did not drain") + + +@pytest.mark.parametrize("mode", ["resident", "offload", "resident_then_offload"]) +def test_suspension_after_capture_resumes_with_the_same_reference(mode, monkeypatch): + victim = SamplingParams(max_tokens=8, ignore_eos=True, priority=1) + arrival = SamplingParams(max_tokens=4, ignore_eos=True) + results = [] + for preempt in (False, True): + e = _ordinary(False, 0.5, scheduler_config=_scheduler()) + seen = _record_modes(e, monkeypatch) + try: + e.add_request("a", [2, 3, 4, 5], victim) + outputs, mark = {}, None + for _ in range(1000): + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + # Stop right after a capture of a generated token: its later + # loops have not consumed the reference yet. + if seen and ("a", CAPTURE) in seen[-1] and e.reference_store.outstanding == 1: + break + assert "a" not in outputs + e.add_request("b", [2, 3], arrival) + if preempt: + waiting = e.scheduler.requests["b"] + assert e.scheduler.preempt(waiting, priority_only=mode != "offload") + saved = e.preemption.snapshots["a"]["references"] + assert saved is not None and e.reference_store.outstanding == 0 + if mode == "resident_then_offload": + assert saved[1].selected.device == e.cache_manager.device + assert e.scheduler.preempt(waiting, waiting_only=True) + saved = e.preemption.snapshots["a"]["references"] + expected = e.cache_manager.device if mode == "resident" else torch.device("cpu") + assert saved[1].selected.device == expected + mark = len(seen) + results.append(_drain(e, outputs)) + assert e.preemption.resumptions == int(preempt) + assert not e.preemption.snapshots + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + if preempt: + # The resumed token consumes its restored reference, not dense. + after = next(plan for plan in seen[mark:] if any(r == "a" for r, _ in plan)) + assert ("a", CONSUME) in after + finally: + e.close() + assert results[0] == results[1] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_kv_pressure_preemption_matches_unpreempted_sparse(asynchronous, monkeypatch): + def run(cache_blocks, preemption): + e = _ordinary( + asynchronous, + 0.5, + cache_config=CacheConfig(cache_blocks, 2, incremental_allocation=True), + scheduler_config=_scheduler(max_num_seqs=3, enable_preemption=preemption), + ) + _bounded_pool(e, monkeypatch) + snapshot = e.model_runner.snapshot_attention_reference + carried = [] + + def counted(request, device): + saved = snapshot(request, device) + carried.append(saved is not None) + return saved + + monkeypatch.setattr(e.model_runner, "snapshot_attention_reference", counted) + try: + outputs = generate(e, count=3) + return outputs, e.preemption.preemptions, sum(carried) + finally: + e.close() + + baseline, none, _ = run(512, False) + pressured, preemptions, carried = run(40, True) + assert none == 0 and preemptions > 0 and carried > 0 + assert pressured == baseline + + +def test_abort_while_suspended_frees_reference_and_id_is_reusable(monkeypatch): + e = _ordinary(False, 0.5, scheduler_config=_scheduler()) + seen = _record_modes(e, monkeypatch) + try: + e.add_request("a", [2, 3, 4], SamplingParams(max_tokens=8, ignore_eos=True, priority=1)) + for _ in range(1000): + e.step() + if seen and ("a", CAPTURE) in seen[-1]: + break + e.add_request("b", [2, 3], SamplingParams(max_tokens=2, ignore_eos=True)) + assert e.scheduler.preempt(e.scheduler.requests["b"], priority_only=True) + assert e.preemption.snapshots["a"]["references"] is not None + e.abort_request("a") + assert "a" not in e.preemption.snapshots + _drain(e, {}) + assert e.reference_store.outstanding == 0 + assert generate(e, count=2) + finally: + e.close() + + +@pytest.mark.parametrize("draft", [1, 2]) +def test_speculative_preemption_stays_rejected(draft): + with pytest.raises(ValueError, match="preemption only for ordinary decode"): + engine( + speculative_config=SpeculativeConfig(2, draft_loops=draft), + scheduler_config=SchedulerConfig(enable_preemption=True), + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("implementation", ["batched", "fused"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_cuda_kv_pressure_preemption_matches_unpreempted(implementation, asynchronous): + from vllm_rlt import SparseDecodeAttentionConfig + + def run(cache_blocks, preemption): + e = _ordinary( + asynchronous, + None, + device="cuda", + cache_config=CacheConfig(cache_blocks, 2, incremental_allocation=True), + scheduler_config=_scheduler(max_num_seqs=3, enable_preemption=preemption), + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=True, key_retention=0.5, implementation=implementation + ), + ) + try: + return generate(e, count=3), e.preemption.preemptions + finally: + e.close() + + baseline, _ = run(512, False) + pressured, preemptions = run(40, True) + torch.cuda.synchronize() + assert preemptions > 0 + assert pressured == baseline + + +@pytest.mark.gpu +@pytest.mark.parametrize("graphs", [False, True]) +@pytest.mark.parametrize("offload", [False, True]) +def test_cuda_pipeline_preemption_restores_each_rows_reference(graphs, offload): + """A token pipeline suspends with rows past their capture loop and resumes them sparse.""" + from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig + + def make(): + return engine( + retention=None, + device="cuda", + speculative_config=SpeculativeConfig(4, draft_loops=1, mode="pipeline"), + scheduler_config=_scheduler(policy="priority"), + execution_config=ExecutionConfig( + async_scheduling=True, + cuda_graphs=graphs, + static_buffers=True, + pad_to_power_of_two=True, + ), + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=True, key_retention=0.5, implementation="fused" + ), + ) + + prompt = [1, 2, 3] + low = SamplingParams(max_tokens=64, ignore_eos=True, priority=10) + e = make() + try: + e.add_request("low", prompt, low) + expected = _drain(e, {})["low"] + finally: + e.close() + + e = make() + try: + e.add_request("low", prompt, low) + request = e.scheduler.requests["low"] + runner = e.speculative_runner + outputs, deep = {}, 0 + + def step(): + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + + # A random tiny model rejects most drafts, so a pipeline often restarts + # from its root. Suspend repeatedly, one more submission later each time + # so the suspension point walks through the pipeline's cycle, until one + # suspension holds a row past its capture loop. + for attempt in range(8): + for _ in range(200): + step() + state = runner.states.get("low") + if state is not None and state.submitted >= 2 + attempt: + break + else: + pytest.fail("the low request did not run pipeline steps") + e.add_request( + f"high{attempt}", [5], SamplingParams(max_tokens=3, ignore_eos=True, priority=0) + ) + for _ in range(200): + step() + if "low" in e.preemption.snapshots: + break + else: + pytest.fail("priority arrival did not suspend the pipeline") + saved = e.preemption.snapshots["low"]["pipeline"] + references, permutation = saved.references + assert sorted(permutation) == list(range(len(references))) + assert "low" not in runner.states + if offload: + e.preemption.suspend(request, offload_kv=True) + saved = e.preemption.snapshots["low"]["pipeline"] + assert saved.references[0].selected.device.type == "cpu" + deep += any(t.depth > 1 for t in saved.tokens) + if deep: + break + assert deep, "no suspension held a row past its capture loop" + assert "low" not in outputs + outputs.update(_drain(e, {})) + assert outputs["low"] == expected + assert e.preemption.resumptions >= 1 + assert not e.preemption.snapshots + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + if graphs: + assert runner.graphs.replays > 0 + finally: + e.close() diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py new file mode 100644 index 00000000..a554dd9b --- /dev/null +++ b/tests/test_sparse_decode_sync.py @@ -0,0 +1,533 @@ +"""Ordinary and synchronous block decode through the shared sparse contract.""" + +from dataclasses import replace + +import pytest +import torch + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.attention import BackendCapabilities, create_backend +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE +from vllm_rlt.engine.stage_engine import StageEngine + + +def engine(*, draft=None, retention=1.0, device="cpu", dtype=torch.float32, **options): + values = dict( + attention_backend="triton" if device == "cuda" else "torch", + cache_config=CacheConfig(512, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=3 + ), + execution_config=ExecutionConfig(static_buffers=True, pad_to_power_of_two=True), + speculative_config=None if draft is None else SpeculativeConfig(2, draft_loops=draft), + sparse_decode_attention_config=( + None + if retention is None + else SparseDecodeAttentionConfig(enabled=True, key_retention=retention) + ), + ) + values.update(options) + return StageEngine(tiny_ouro_model(device, dtype), **values) + + +def generate(e, count=3): + for i in range(count): + e.add_request( + str(i), [2, 3, 4, 5][: 2 + i % 3], SamplingParams(max_tokens=8, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not e.has_unfinished_requests(): + break + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + assert len(outputs) == count + assert all(len(tokens) == 8 for tokens in outputs.values()) + assert e.cache_manager.num_used_blocks == 0 + if e.reference_store is not None: + assert e.reference_store.outstanding == 0 + return outputs + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_full_retention_matches_dense_and_reclaims_slots(draft): + dense = engine(draft=draft, retention=None) + sparse = engine(draft=draft) + try: + assert generate(sparse) == generate(dense) + finally: + sparse.close() + dense.close() + + +def _ordinary(asynchronous, retention, device="cpu", dtype=torch.float32, **options): + return engine( + retention=retention, + device=device, + dtype=dtype, + exit_config=ExitConfig("ouro_delayed"), + execution_config=ExecutionConfig( + async_scheduling=asynchronous, static_buffers=True, pad_to_power_of_two=True + ), + **options, + ) + + +def _bounded_pool(e, monkeypatch): + """Fail if live references ever exceed one per scheduled request.""" + store = e.reference_store + reserve = store.reserve + + def checked(*args, **kwargs): + lease = reserve(*args, **kwargs) + assert store.outstanding <= e.scheduler.config.max_num_seqs + return lease + + monkeypatch.setattr(store, "reserve", checked) + + +@pytest.mark.parametrize("retention", [0.5, 1.0]) +def test_async_ordinary_decode_matches_synchronous(retention, monkeypatch): + sync = _ordinary(False, retention) + asynchronous = _ordinary(True, retention) + _bounded_pool(asynchronous, monkeypatch) + try: + # Staggered prompt lengths keep requests at different depths per batch. + assert generate(asynchronous, count=4) == generate(sync, count=4) + finally: + sync.close() + asynchronous.close() + + +def test_async_ordinary_decode_reclaims_after_abort_and_reuse(monkeypatch): + e = _ordinary(True, 0.5) + _bounded_pool(e, monkeypatch) + try: + e.add_request("a", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + assert e.reference_store.outstanding == 1 + e.abort_request("a") + for _ in range(100): + if not e.has_unfinished_requests(): + break + e.step() + assert e.reference_store.outstanding == 0 + # The same id restarts with a fresh incarnation. + assert generate(e, count=2) + finally: + e.close() + + +@pytest.mark.gpu +@pytest.mark.parametrize("implementation", ["batched", "fused"]) +def test_cuda_async_ordinary_decode_matches_synchronous(implementation, monkeypatch): + sparse = SparseDecodeAttentionConfig( + enabled=True, key_retention=0.5, implementation=implementation + ) + options = dict(device="cuda", dtype=torch.bfloat16, sparse_decode_attention_config=sparse) + sync = _ordinary(False, 0.5, **options) + asynchronous = _ordinary(True, 0.5, **options) + _bounded_pool(asynchronous, monkeypatch) + store = asynchronous.reference_store + host_count = store.host_count + published = [] + + def checked(slot, layer): + # Every consume is sized from a count the host published at capture. + count = host_count(slot, layer) + assert count is not None + # ...and it is the count the capture actually wrote. + assert int(store.counts[slot, layer].max()) == count + published.append((slot, layer, count)) + return count + + monkeypatch.setattr(store, "host_count", checked) + try: + assert generate(asynchronous, count=4) == generate(sync, count=4) + assert published + finally: + sync.close() + asynchronous.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_plans_follow_tokens_and_keep_prompt_dense(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + seen = [] + original = e.cache_manager._attend_prepared + + def attend(layer, kv, q): + plan = getattr(kv, "reference_plan", None) + if plan is not None: + assert len(plan.modes) == len(q) + for row, mode in enumerate(plan.modes): + if mode == DENSE: + continue + allocation, depth, position = kv.rows[row] + request_id = next(rid for rid, owner in kv.allocations if owner is allocation) + assert position >= len(e.scheduler.requests[request_id].prompt_token_ids) + assert mode == (CAPTURE if depth == 1 else CONSUME) + assert int(e.reference_store.position[plan.slots[row]]) == position + seen.append((mode, position, plan.slots[row])) + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", attend) + try: + generate(e) + assert {mode for mode, _, _ in seen} == {CAPTURE, CONSUME} + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cancel_and_request_id_reuse_cannot_reuse_old_reference(draft): + e = engine(draft=draft, retention=0.5) + try: + for _ in range(2): + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=32, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding or ( + draft is not None and e.speculative_runner.stats.rounds + ): + break + e.abort_request("r") + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + finally: + e.close() + + +def test_context_bound_is_enforced_before_request_admission(): + e = engine( + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=True, max_context_tokens=8 + ) + ) + try: + with pytest.raises(ValueError, match="sparse.*context"): + e.add_request("too-long", [2, 3, 4], SamplingParams(max_tokens=8)) + assert "too-long" not in e.scheduler.requests + finally: + e.close() + + +@pytest.mark.parametrize("setting", ["graphs", "pipeline", "preemption", "shared"]) +def test_unqualified_execution_modes_fail_explicitly(setting): + options = { + "graphs": dict(execution_config=ExecutionConfig(cuda_graphs=True)), + "pipeline": dict(speculative_config=SpeculativeConfig(2, mode="pipeline")), + "preemption": dict( + speculative_config=SpeculativeConfig(2), + scheduler_config=SchedulerConfig(enable_preemption=True), + ), + "shared": dict(cache_config=CacheConfig(512, 2, layout="shared")), + }[setting] + with pytest.raises(ValueError, match="sparse decode attention"): + engine(**options) + + +def test_backend_explicitly_declares_sparse_reference_support(): + selected = create_backend( + "auto", + torch.device("cpu"), + torch.float32, + 8, + 2, + required_capabilities=BackendCapabilities(sparse_decode=True), + ) + assert selected.capabilities.sparse_decode + + +@pytest.mark.gpu +@pytest.mark.parametrize("kv_heads", [2, 16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): + from vllm_rlt.core.kv_cache_manager import KVCacheManager + from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + # BF16 rounding of Torch softmax and Triton split-KV can differ even for + # identical operands. Dense/capture rows and the full-retention control + # must preserve the configured backend, including in a mixed plan. + torch.manual_seed(157) + rows, heads, dim, length = 32, 16, 128, 130 + cache = KVCacheManager( + 1, + kv_heads, + dim, + rows * 9, + 16, + 1, + device="cuda", + dtype=torch.bfloat16, + backend="triton", + ) + store = AttentionReferenceStore( + slots=rows, layers=1, heads=heads, head_dim=dim, key_bucket=length, device="cuda" + ) + ids = [str(i) for i in range(rows)] + owners = [] + for rid in ids: + assert cache.allocate(rid, length) + keys = torch.randn(length, kv_heads, dim, device="cuda", dtype=torch.bfloat16) + values = torch.randn_like(keys) + cache.write(0, [rid] * length, [0] * length, list(range(length)), keys, values) + lease = store.reserve(rid, 0, 1, generation=1) + owners.append((lease, 0)) + batch = cache._prepare_batch(ids, [0] * rows, [length - 1] * rows) + q = torch.randn(rows, heads, dim, device="cuda", dtype=torch.bfloat16) + expected = cache._attend_prepared(0, batch, q) + modes = [DENSE] + [CAPTURE] * (rows - 1) + plan = store.build_plan(owners, modes, key_retention=retention) + actual = cache._attend_prepared(0, replace(batch, reference_plan=plan), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + for row in range(1, rows): + torch.testing.assert_close( + store.output[plan.slots[row], 0], expected[row].float(), atol=0, rtol=0 + ) + if retention == 1.0: + # Changed later-loop Q/K/V prevents an identity-only correction test. + q = torch.randn_like(q) + cache._write_prepared( + 0, batch, torch.randn_like(keys[:rows]), torch.randn_like(keys[:rows]) + ) + expected = cache._attend_prepared(0, batch, q) + consume = store.build_plan( + owners, [DENSE] + [CONSUME] * (rows - 1), key_retention=retention + ) + actual = cache._attend_prepared(0, replace(batch, reference_plan=consume), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert int(store.uses.sum()) == rows - 1 + for lease, _ in owners: + store.release(lease) + for rid in ids: + cache.free(rid) + torch.cuda.synchronize() + assert store.outstanding == cache.num_used_blocks == 0 + + +@pytest.mark.parametrize("draft", [1, 2]) +def test_sparse_verification_matches_serial_same_input_hidden_logits_and_kv(draft): + from tests.helpers import make_cpu_cache + from tests.test_sparse_decode_attention import make_reference_store + from vllm_rlt.attention.sparse import reference_mode + from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput + from vllm_rlt.request import Request, RequestPhase + + e = engine(draft=draft, retention=0.5) + model, cache = e.model, e.cache_manager + prefix, current = [3, 4, 5], 7 + assert cache.allocate("r", 16) + for position, token in enumerate(prefix): + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + hidden, _ = model.recurrent(hidden, ["r"], [depth], [position], cache) + request = Request("r", prefix, SamplingParams(ignore_eos=True), generated_token_ids=[current]) + seen, heads = {}, [] + recurrent = model.recurrent_prepared + + def record(hidden, kv, cache, **kwargs): + out, gate = recurrent(hidden, kv, cache, **kwargs) + for row, (_, depth, position) in enumerate(kv.rows): + seen[depth, position] = out[row].clone() + return out, gate + + model.recurrent_prepared = record + hook = model.lm_head.register_forward_hook( + lambda module, inputs, output: heads.append(output.detach().clone()) + ) + try: + e.speculative_runner.submit( + SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(request, 3, 3)]) + ).collect() + finally: + hook.remove() + model.recurrent_prepared = recurrent + assert e.reference_store.outstanding == 0 + tokens = prefix + [current] + tokens.extend(int(model.coda(seen[draft - 1, pos]).argmax()) for pos in (3, 4)) + oracle = make_cpu_cache(model.config) + assert oracle.allocate("r", 16) + store = make_reference_store(model) + for position, token in enumerate(tokens): + lease = store.reserve("r", 0, 1, generation=position) + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + kv = oracle._prepare_batch(["r"], [depth], [position]) + mode = reference_mode(depth) + if position >= len(prefix) and mode != DENSE: + plan = store.build_plan([(lease, 0)], [mode], key_retention=0.5) + kv = replace(kv, reference_plan=plan) + hidden, _ = model.recurrent_prepared(hidden, kv, oracle, compute_gate=False) + if position >= len(prefix): + torch.testing.assert_close(seen[depth, position], hidden[0], atol=4e-5, rtol=4e-5) + if position >= len(prefix): + torch.testing.assert_close( + heads[-1][position - len(prefix)], model.coda(hidden)[0], atol=4e-5, rtol=4e-5 + ) + store.release(lease) + for depth in range(4): + for layer in range(model.config.num_hidden_layers): + for actual, expected in zip( + cache.read(layer, "r", depth, len(tokens)), + oracle.read(layer, "r", depth, len(tokens)), + ): + torch.testing.assert_close(actual, expected, atol=4e-5, rtol=4e-5) + cache.free("r") + oracle.free("r") + e.close() + + +@pytest.mark.parametrize("draft", [1, 2]) +@pytest.mark.parametrize("accepted", [0, 1, 2]) +def test_sync_rejection_retires_all_references_before_replacement(draft, accepted, monkeypatch): + e = engine(draft=draft, retention=0.5) + try: + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=16, ignore_eos=True)) + while not e.step(): + pass + request = e.scheduler.requests["r"] + start = request.position + original = e.speculative_runner._coda + calls = 0 + + def controlled(hidden): + nonlocal calls + logits = torch.full((len(hidden), 64), -100.0) + if calls < 2: + logits[:, 10 + calls] = 100 + else: + for row in range(3): + logits[row, 10 + row if row < accepted else 20 + row] = 100 + calls += 1 + return logits + + seen = [] + build = e.reference_store.build_plan + + def record(*args, **kwargs): + plan = build(*args, **kwargs) + seen.append(plan) + return plan + + monkeypatch.setattr(e.reference_store, "build_plan", record) + monkeypatch.setattr(e.speculative_runner, "_coda", controlled) + output = e.step()[0] + expected = [10 + row if row < accepted else 20 + row for row in range(accepted + 1)] + assert output.token_ids[-(accepted + 1) :] == expected + assert request.position == start + accepted + 1 + assert e.reference_store.outstanding == 0 + old = seen[0] + with pytest.raises(RuntimeError, match="stale"): + e.reference_store.read_reference(old, 0, 0) + monkeypatch.setattr(e.speculative_runner, "_coda", original) + e.step() + assert seen[-1].incarnations[0] > old.incarnations[0] + assert e.speculative_runner.stats.accepted_tokens >= accepted + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_failed_attention_releases_lease_and_allows_fresh_request(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + original = e.cache_manager._attend_prepared + + def fail(layer, kv, q): + if getattr(kv, "reference_plan", None) is not None: + raise RuntimeError("injected attention failure") + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", fail) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + with pytest.raises(RuntimeError, match="injected"): + for _ in range(100): + e.step() + assert e.reference_store.outstanding == 0 + assert "r" not in e.scheduler.requests + assert e.cache_manager.num_used_blocks == 0 + monkeypatch.setattr(e.cache_manager, "_attend_prepared", original) + generate(e, count=1) + finally: + e.close() + + +@pytest.mark.gpu +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cuda_sparse_sync_executes_and_drains(draft): + e = engine(draft=draft, retention=0.5, device="cuda", dtype=torch.bfloat16) + try: + generate(e) + torch.cuda.synchronize() + assert not bool(e.reference_store.counts.any()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_ordinary_retirement_survives_boundary_allocator_reuse(monkeypatch): + e = engine(device="cuda", dtype=torch.bfloat16) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + store = e.reference_store + assert store.outstanding == 1 + slot = e.model_runner.attention_references["r"][2].slots[0] + assert slot != 0 + torch.cuda.synchronize() + store.counts[0] = 77 + torch.cuda.synchronize() + clear = store._clear + pointers = [] + + def record(indices): + assert torch.cuda.current_stream() == e.model_runner.core_stream + pointers.append(indices.data_ptr()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + # Hold core while boundary allocations try to reuse any released index + # storage. Correct stream ownership prevents that reuse before clearing. + with torch.cuda.stream(e.model_runner.core_stream): + torch.cuda._sleep(20_000_000) + with torch.cuda.stream(e.model_runner.boundary_stream): + e.model_runner._release_attention_reference("r") + poison = [torch.zeros(1, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert pointers and len(poison) == 32 + assert not bool(store.counts[slot].any()) + assert bool((store.counts[0] == 77).all()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_close_drains_retirement_of_an_active_reference(): + e = engine(device="cuda", dtype=torch.bfloat16) + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + assert e.reference_store.outstanding == 1 + e.close() + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + assert e.model_runner.core_stream.query() diff --git a/vllm_rlt/__init__.py b/vllm_rlt/__init__.py index d0b346bc..c813eec0 100644 --- a/vllm_rlt/__init__.py +++ b/vllm_rlt/__init__.py @@ -5,6 +5,7 @@ ExecutionConfig, ExitConfig, SchedulerConfig, + SparseDecodeAttentionConfig, SpeculativeConfig, ) from vllm_rlt.entrypoints.llm import LLM @@ -19,6 +20,7 @@ "ExecutionConfig", "ExitConfig", "SchedulerConfig", + "SparseDecodeAttentionConfig", "SpeculativeConfig", "SamplingParams", "RequestOutput", diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..045c1a23 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -23,6 +23,8 @@ class BackendCapabilities: packed_matches_decode: bool = False # Consume GPU-owned context lengths without a host read between submissions. device_context_lengths: bool = False + # Shared paged sparse-reference capture/correction dispatch is available. + sparse_decode: bool = False @dataclass @@ -45,9 +47,9 @@ def generation(self): return getattr(self.kernel, "generation", None) -_TORCH_CAPABILITIES = BackendCapabilities() +_TORCH_CAPABILITIES = BackendCapabilities(sparse_decode=True) _TRITON_CAPABILITIES = BackendCapabilities( - cuda_graphs=True, async_scheduling=True, device_context_lengths=True + cuda_graphs=True, async_scheduling=True, device_context_lengths=True, sparse_decode=True ) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True), @@ -128,6 +130,7 @@ def create_backend( "cuda_graphs", "async_scheduling", "device_context_lengths", + "sparse_decode", ) if getattr(required_capabilities, name) and not getattr(implementation.capabilities, name) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py new file mode 100644 index 00000000..668b25a5 --- /dev/null +++ b/vllm_rlt/attention/sparse.py @@ -0,0 +1,677 @@ +"""Sparse decode attention arithmetic and its per-row plan (RFC #156, section 3.2). + +The sparse decode algorithm keeps the attention output of a reference loop and +replaces only the contribution of the retained keys in a later loop. The +reference loop is loop 2 in the RFC's one-based numbering, which is depth index +1 at runtime. + +This module holds the arithmetic and the plan that describes one submission: + +``reference_statistics`` + the dense reference attention of one row, plus the retained state + ``(S, m, O_ref, C_ref)``, +``sparse_correction`` + the corrected output of a later loop, ``O = O_ref - C_ref + m * C_new``, +``SparseDecodePlan`` + the row modes and reference slots of one submission. + +The arithmetic is deliberately plain PyTorch. The retained state is small, so a +straightforward implementation is easier to check against the worked example in +the RFC than an optimized kernel, and it keeps the contract that the attention +kernels must satisfy visible in one place. + +The plan is the named view between the runner and the attention layer. The +runner owns the reference slots and their lifetime; the attention layer only +borrows the plan for the rows of one submission. Row modes and slots are +host-known: the runner decides both before it submits, so a submission never +depends on a device readback. +""" + +import math +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +from vllm_rlt.kernels.sparse_capture import TRITON_AVAILABLE, fused_capture_scores +from vllm_rlt.kernels.sparse_consume import selected_scores as fused_selected_scores +from vllm_rlt.kernels.sparse_consume import selected_values as fused_selected_values + +# Row modes for one recurrent submission. +DENSE = 0 +CAPTURE = 1 +CONSUME = 2 + +# The reference loop is loop 2 in the RFC's one-based numbering, which is +# zero-based depth index 1 at runtime. +REFERENCE_DEPTH = 1 + + +def reference_mode(depth: int) -> int: + """Row mode of a recurrent step at the given zero-based loop depth.""" + if depth == REFERENCE_DEPTH: + return CAPTURE + if depth > REFERENCE_DEPTH: + return CONSUME + return DENSE + + +@dataclass(frozen=True) +class SparseDecodePlan: + """The reference plan of one recurrent submission. + + ``modes`` and ``slots`` are host-known: the runner decides both before it + submits, so a submission never depends on a device readback. ``slots`` holds + one reference slot per row, or ``-1`` for a row that owns no reference. + ``incarnations`` identifies the token that occupies each slot; a plan built + for a superseded incarnation is stale and must be rejected. + + A dense row runs the ordinary attention path. A capture row is the + reference loop: it computes the dense attention and publishes the retained + state. A consume row is a later loop: it corrects the reference output with + the retained keys. + + ``store`` is the runner-owned reference pool, borrowed by the attention + layer through duck typing so that model code needs no import of the worker + package. + """ + + modes: tuple[int, ...] + slots: tuple[int, ...] + incarnations: tuple[int, ...] + key_retention: float + key_bucket: int | None + slot_ids: torch.Tensor | None = None + store: object | None = None + + def groups(self) -> dict[int, list[int]]: + """Rows per mode, in submission order.""" + grouped: dict[int, list[int]] = {} + for row, mode in enumerate(self.modes): + grouped.setdefault(mode, []).append(row) + return {mode: rows for mode, rows in grouped.items() if rows} + + @property + def reference_depth(self) -> int: + """The zero-based depth that captures the reference state.""" + return REFERENCE_DEPTH + + def dense(self) -> bool: + """Whether the submission has no sparse work at all.""" + return all(mode == DENSE for mode in self.modes) + + +@dataclass(frozen=True) +class DeviceSparseDecodePlan: + """A reference plan whose row modes and slots are decided on the device. + + Token pipelines advance, shift and restart their rows on the GPU, so the + host knows neither a row's depth nor which of its request's slots it owns. + ``modes`` and ``slot_ids`` are device tensors with one entry per row; + ``leases`` own every slot a row can name and are checked live on the host. + Only the fused kernels read such a plan. + """ + + modes: torch.Tensor + slot_ids: torch.Tensor + leases: tuple + key_retention: float + key_bucket: int | None + store: object | None = None + # False when no row can be in its capture loop (the host knows which loops a + # step runs); the fused path then launches no selection kernels. + may_capture: bool = True + + def dense(self) -> bool: + """Unknown on the host: the device may route any row to sparse work.""" + return False + + +def retained_count(length: int, key_retention: float) -> int: + """How many keys a query with ``length`` legal keys retains. + + Follows the FlashLoop reference (``aeaadee``, ``select_decode_keys``): + ``k = min(length, max(1, ceil(length * key_retention)))``. The count is a + pure function of the retention fraction, so ``key_retention=1.0`` stays an + exact full-retention control; index-buffer capacity lives in the store and + never changes the algorithm. + """ + if key_retention <= 0 or key_retention > 1: + raise ValueError("key_retention must be in (0, 1]") + if length < 1: + return 0 + return min(length, max(1, math.ceil(length * key_retention))) + + +def expand_kv(tensor: torch.Tensor, groups: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention.""" + if groups == 1: + return tensor + return tensor.repeat_interleave(groups, dim=1) + + +def select_keys(probabilities: torch.Tensor, count: int) -> torch.Tensor: + """Pick the retained key positions of every head. + + ``probabilities`` is ``[heads, length]``. Equal probabilities are broken by + the smaller absolute position, so the selection does not depend on the order + of the rows in a batch. + """ + if count < 1: + raise ValueError("a selection needs at least one key") + if count > probabilities.shape[1]: + raise ValueError("cannot retain more keys than the query has") + # A stable descending sort keeps equal probabilities in position order. + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + return order[:, :count].sort(dim=-1).values + + +def reference_statistics( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + *, + key_retention: float, + key_bucket: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Dense reference attention and the retained state of one row. + + ``q`` is ``[heads, dim]``, ``keys`` and ``values`` are + ``[length, kv_heads, dim]``. Returns ``(output, contribution, mass, + positions)`` where output and contribution are ``[heads, dim]``, mass is + ``[heads]`` and positions are the retained absolute key positions. + """ + heads, dim = q.shape + length = keys.shape[0] + groups = heads // keys.shape[1] + scale = 1.0 / dim**0.5 + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q.float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output = torch.einsum("hl,lhd->hd", probabilities, expanded_values) + count = retained_count(length, key_retention) + # Buffer capacity never truncates the algorithm: a count above it fails + # loudly instead of silently degrading retention. + if key_bucket is not None and count > key_bucket: + raise ValueError(f"reference capacity {key_bucket} cannot hold {count} retained keys") + positions = select_keys(probabilities, count) + gathered = probabilities.gather(1, positions) + # Index the per-head value rows: positions is [heads, count], so the head + # index broadcasts against it and selects each head's own key rows. + head_index = torch.arange(heads, device=q.device)[:, None] + selected_values = expanded_values[positions, head_index] + contribution = torch.einsum("hk,hkd->hd", gathered, selected_values) + mass = gathered.sum(dim=-1) + return output, contribution, mass, positions + + +def sparse_correction( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + mass: torch.Tensor, + output: torch.Tensor, + contribution: torch.Tensor, +) -> torch.Tensor: + """Correct the reference output with the retained keys of a later loop. + + ``keys`` and ``values`` are ``[heads, retained, dim]`` and already carry the + keys of the later loop. ``mass``, ``output`` and ``contribution`` are the + retained reference state. + """ + dim = q.shape[-1] + scale = 1.0 / dim**0.5 + scores = torch.einsum("hd,hkd->hk", q.float(), keys.float()) * scale + probabilities = torch.softmax(scores, dim=-1) + new_contribution = torch.einsum("hk,hkd->hd", probabilities, values.float()) + return output - contribution + mass.float().unsqueeze(-1) * new_contribution + + +def gather_retained( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + selected: torch.Tensor, + block_size: int, + groups: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather each query head's retained keys and values from the paged cache. + + ``selected`` holds absolute key positions per query head, so every head + reads its own retained set. Returns ``[heads, retained, dim]`` tensors. + """ + index = selected.long() + blocks = block_table[index // block_size] + offsets = index % block_size + kv_index = (torch.arange(index.shape[0], device=index.device) // groups)[:, None] + kv_index = kv_index.expand_as(index) + return key_cache[blocks, offsets, kv_index], value_cache[blocks, offsets, kv_index] + + +def reference_paged_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + plan: SparseDecodePlan, + *, + layer: int, + dense_attention: Callable | None = None, +) -> torch.Tensor: + """Reference sparse decode attention over a paged KV cache (RFC #156 3.2). + + This is the numerical oracle for the sparse decode column, not an optimized + kernel: it walks rows on the host, reads the same paged tensors the dense + path reads, and publishes or consumes the retained state through the plan's + store. Dense rows reproduce the ordinary attention path exactly. + + ``key_cache``/``value_cache`` are one layer of + ``[physical block, token, KV head, dimension]``; ``block_tables`` and + ``context_lengths`` describe the submission rows exactly as the dense + backend receives them. When supplied, ``dense_attention`` preserves that + backend's outputs for dense/capture rows and the full-retention control. + Omitting it forces the independent Torch arithmetic oracle, including the + full-retention correction; that path need not be bitwise equal to Triton. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + block_size = key_cache.shape[1] + heads = q.shape[1] + groups = heads // key_cache.shape[2] + scale = 1.0 / math.sqrt(q.shape[-1]) + output = torch.empty_like(q) + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. Only the rows needing dense output use this result. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + captured = consumed = False + for row, length in enumerate(context_lengths.tolist()): + if length == 0: + output[row].zero_() + continue + positions = torch.arange(length, device=q.device) + blocks = block_tables[row, positions // block_size] + offsets = positions % block_size + mode = plan.modes[row] + if mode == CONSUME: + selected, mass, reference, contribution = store.read_reference(plan, layer, row) + if dense_output is not None and plan.key_retention == 1.0: + output[row] = dense_output[row] + consumed = True + continue + keys, values = gather_retained( + key_cache, value_cache, block_tables[row], selected, block_size, groups + ) + output[row] = sparse_correction(q[row], keys, values, mass, reference, contribution).to( + q.dtype + ) + consumed = True + continue + keys = key_cache[blocks, offsets] + values = value_cache[blocks, offsets] + if mode == CAPTURE: + reference, contribution, mass, selected = reference_statistics( + q[row], + keys, + values, + key_retention=plan.key_retention, + key_bucket=store.key_bucket, + ) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence. Keep + # selector probabilities/contributions in the Torch FP32 oracle. + reference = dense_output[row].float() + output[row] = reference.to(q.dtype) + store.write_reference( + plan, + layer, + row, + selected=selected, + counts=torch.full((heads,), selected.shape[-1], dtype=torch.int32, device=q.device), + mass=mass, + output=reference, + contribution=contribution, + ) + captured = True + continue + # Dense row: the ordinary attention path. + if dense_output is not None: + output[row] = dense_output[row] + continue + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q[row].float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output[row] = torch.einsum("hl,lhd->hd", probabilities, expanded_values).to(q.dtype) + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output + + +def sparse_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, +): + """Use bounded row batches on CUDA, retaining the scalar numerical oracle. + + A store configured for the ``fused`` implementation routes uniform capture + or consume submissions, and every device plan, to the fused kernels. Other + CUDA submissions keep + the batched path, whose capture rows take the fused paged score pass when a + dense backend is available; ``batched_paged_attention`` without + ``fused_scores`` still reproduces the scalar oracle bit for bit. + """ + if isinstance(plan, DeviceSparseDecodePlan): + from vllm_rlt.kernels import sparse_decode + + if not sparse_decode.supports(plan, q, dense_attention): + raise ValueError("a device sparse plan needs the fused CUDA kernels") + return sparse_decode.fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + if not q.is_cuda: + return reference_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + if getattr(plan.store, "implementation", "batched") == "fused": + from vllm_rlt.kernels import sparse_decode + + if sparse_decode.supports(plan, q, dense_attention): + return sparse_decode.fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + return batched_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + fused_scores=TRITON_AVAILABLE and dense_attention is not None, + ) + + +def _rowwise_product(equation, left, right): + """Keep each row's cuBLAS problem shape equal to the scalar oracle. + + Folding rows into the head dimension can choose a different FP32 reduction + kernel. Its small error can cross a BF16 rounding boundary and change a + greedy token, so batch the gathers/selection but retain row-wise products. + """ + # Prepare the exact views einsum uses once for the group, instead of + # repeating its permutation/reshape bookkeeping for every row. + if equation == "hd,lhd->hl": + right = right.permute(0, 2, 3, 1) + elif equation == "hl,lhd->hd": + right = right.transpose(1, 2) + elif equation == "hd,hkd->hk": + right = right.transpose(-1, -2) + elif equation != "hk,hkd->hd": + raise ValueError(f"unsupported reference product {equation}") + return torch.stack([torch.bmm(a, b).squeeze(1) for a, b in zip(left.unsqueeze(-2), right)]) + + +def batched_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, + fused_scores: bool = False, +): + """Batch equal-length rows without padding or changing retained-key selection. + + Each head still uses the same FP32 matrix products, stable sort and softmax + as the scalar oracle. Grouping exact lengths avoids padded reductions and + bounds gather scratch independently of the reference pool's key capacity. + There is one count readback per layer, rather than one per consume row. + + ``dense_attention`` keeps the numerical contract of the scalar path: dense + and capture rows publish the selected backend's output, the actually + rounded recurrence input becomes ``O_ref``, and a full-retention consume row + returns that same backend output as the control. Selector probabilities and + the retained contribution stay in the Torch FP32 oracle. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + store._require_live(plan) + lengths = context_lengths.tolist() + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + consume_rows = [r for r, mode in enumerate(plan.modes) if mode == CONSUME and lengths[r]] + counts = {} + if consume_rows: + # The retained count is host-known from the capture. Reading it back + # from the device would sync every layer, which RFC #156 section 6.2 + # forbids. + counts = {row: store.host_count(plan.slots[row], layer) for row in consume_rows} + missing = [row for row, count in counts.items() if count is None] + if missing: + # The fused capture decides its count on the device (graph replays + # change lengths without the host), so a mixed submission that + # falls back here after it reads those counts back once. + slots = plan.slot_ids[missing] + counts.update(zip(missing, store.counts[slots, layer].amax(dim=-1).tolist())) + if any(count == 0 for count in counts.values()): + raise RuntimeError("a consume row has no captured reference") + grouped = {} + for row, length in enumerate(lengths): + if length: + mode = plan.modes[row] + extent = counts[row] if mode == CONSUME else length + grouped.setdefault((mode, extent), []).append(row) + output = torch.zeros_like(q) + heads, dim = q.shape[1:] + groups = heads // key_cache.shape[2] + block_size = key_cache.shape[1] + scale = 1.0 / math.sqrt(dim) + head_index = torch.arange(heads, device=q.device)[None, :, None] + kv_index = head_index // groups + captured = consumed = False + for (mode, extent), all_rows in grouped.items(): + # The fused path materialises only the score matrix, so its budget is + # the scores themselves; the expanded path budgets an FP32 K row. Both + # must allow the whole group at once, or a 32k row degenerates into one + # launch per row. + if fused_scores: + chunk = max(1, min(32, 512 * 1024 * 1024 // (heads * extent * 4))) + else: + chunk = max(1, min(32, 64 * 1024 * 1024 // (extent * heads * dim))) + for start in range(0, len(all_rows), chunk): + rows = all_rows[start : start + chunk] + # Contiguous rows come from a device-side arange; otherwise a pinned + # source keeps the copy from draining the stream. + if rows[0] + len(rows) == rows[-1] + 1: + row_ids = torch.arange(rows[0], rows[0] + len(rows), device=q.device) + else: + row_ids = torch.tensor(rows, dtype=torch.int64, pin_memory=True).to( + q.device, non_blocking=True + ) + if mode == DENSE and dense_output is not None: + output[row_ids] = dense_output[row_ids] + continue + if mode == CONSUME and dense_output is not None and plan.key_retention == 1.0: + output[row_ids] = dense_output[row_ids] + consumed = True + continue + slots = plan.slot_ids[row_ids] + chunk_slots = [plan.slots[row] for row in rows] + queries = q[row_ids].float() + tables = block_tables[row_ids] + batch_index = torch.arange(len(rows), device=q.device)[:, None, None] + if mode == CONSUME: + if dense_output is not None and plan.key_retention == 1.0: + # Full retention is the control: same backend as dense. + output[row_ids] = dense_output[row_ids] + consumed = True + continue + # Contiguous slots turn these gathers into views, which is the + # difference between one index kernel per tensor and none. + if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: + span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + selected = store.selected[span, layer, :, :extent].to(torch.int32) + ref_output = store.output[span, layer] + ref_contribution = store.contribution[span, layer] + ref_mass = store.mass[span, layer] + else: + columns = torch.arange(extent, device=q.device)[None, None, :] + selected = store.selected[slots[:, None, None], layer, head_index, columns] + ref_output = store.output[slots, layer] + ref_contribution = store.contribution[slots, layer] + ref_mass = store.mass[slots, layer] + if fused_scores: + # Walk the retained indices in place: no dense K/V gather + # and no per-row cuBLAS. + scores = ( + fused_selected_scores( + q[row_ids], + key_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + * scale + ) + probabilities = torch.softmax(scores, dim=-1) + new = fused_selected_values( + probabilities, + value_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + else: + selected = selected.long() + blocks = tables[batch_index, selected // block_size] + offsets = selected % block_size + keys = key_cache[blocks, offsets, kv_index].float() + values = value_cache[blocks, offsets, kv_index].float() + scores = _rowwise_product("hd,hkd->hk", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + new = _rowwise_product("hk,hkd->hd", probabilities, values) + corrected = ref_output - ref_contribution + ref_mass.unsqueeze(-1) * new + output[row_ids] = corrected.to(q.dtype) + consumed = True + continue + positions = torch.arange(extent, device=q.device) + blocks = tables[:, positions // block_size] + offsets = positions % block_size + if fused_scores and dense_output is not None: + # One pass over the paged cache: no expanded FP32 K row, and + # no per-row cuBLAS. The dot product is still FP32; only the + # reduction order differs from the oracle. The context bound is + # host-known here, so the kernel never reads it back. + chunk_lengths = [lengths[row] for row in rows] + chunk_context = max(chunk_lengths) + scores = fused_capture_scores( + q[row_ids], + key_cache, + tables, + context_lengths[row_ids], + layer=layer, + block_size=block_size, + context=chunk_context, + ragged=min(chunk_lengths) != chunk_context, + ) + scores = scores * scale + else: + keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence; the + # selector state below stays in the Torch FP32 oracle. + reference = dense_output[row_ids].float() + else: + values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + reference = _rowwise_product("hl,lhd->hd", probabilities, values) + output[row_ids] = reference.to(q.dtype) + if mode == CAPTURE: + count = retained_count(extent, plan.key_retention) + store.ensure_key_capacity(count) + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + selected = order[..., :count].sort(dim=-1).values + gathered = probabilities.gather(-1, selected) + # Gather only retained V directly from the paged cache. The + # FP32 values and reduction geometry are identical to gathering + # from the full GQA-expanded tensor, without materializing it. + selected_blocks = tables[batch_index, selected // block_size] + selected_values = value_cache[ + selected_blocks, selected % block_size, kv_index + ].float() + contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) + # Only the live prefix is written, never the whole capacity. + if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: + span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + store.selected[span, layer, :, :count] = selected.to(torch.int32) + store.counts[span, layer] = count + store.mass[span, layer] = gathered.sum(dim=-1) + store.output[span, layer] = reference + store.contribution[span, layer] = contribution + else: + columns = torch.arange(count, device=q.device)[None, None, :] + store.selected[slots[:, None, None], layer, head_index, columns] = selected.to( + torch.int32 + ) + store.counts[slots, layer] = count + store.mass[slots, layer] = gathered.sum(dim=-1) + store.output[slots, layer] = reference + store.contribution[slots, layer] = contribution + for row in rows: + store.note_count(plan.slots[row], layer, count) + captured = True + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 648781d7..d02c7c57 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -190,3 +190,41 @@ 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") + + +@dataclass(frozen=True) +class SparseDecodeAttentionConfig: + """Sparse decode attention group (RFC #156, section 7.2). + + ``enabled`` recomputes attention contributions only for the retained keys + of the loop-2 reference; ``key_retention`` is the retained key fraction per + query and head relative to that query's causal key count, where ``1.0`` + provides the full-retention control. ``key_bucket`` is only the reference + index-buffer capacity, never an algorithm knob: the retained count follows + ``key_retention`` alone, and ``None`` sizes the buffer from the model + context and the retention fraction. + """ + + enabled: bool = False + key_retention: float = 1.0 + key_bucket: int | None = None + # Supported serving context bound for the reference index buffer. ``None`` + # falls back to the model context; set it to the real serving bound so the + # pool is not sized for a context the deployment never uses. + max_context_tokens: int | None = None + # ``batched`` keeps the scalar oracle's reduction shapes (bitwise contract); + # ``fused`` launches fixed per-layer kernels without host readbacks and + # changes consume-row numerics (see vllm_rlt.kernels.sparse_decode). + implementation: str = "batched" + + def __post_init__(self): + if type(self.enabled) is not bool: + raise ValueError("enabled must be a boolean") + if self.implementation not in ("batched", "fused"): + raise ValueError("implementation must be 'batched' or 'fused'") + if not math.isfinite(self.key_retention) or not 0 < self.key_retention <= 1: + raise ValueError("key_retention must be in (0, 1]") + if self.key_bucket is not None: + _positive("key_bucket", self.key_bucket) + if self.max_context_tokens is not None: + _positive("max_context_tokens", self.max_context_tokens) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..a6595f28 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -23,6 +23,7 @@ create_backend, validate_backend_name, ) +from vllm_rlt.attention.sparse import sparse_paged_attention from vllm_rlt.request import ExecutionBatch from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder @@ -718,6 +719,21 @@ def _attend_prepared( self._validate_layer(layer) self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) + if not isinstance(batch, DeviceKVMetadata): + for allocation, depth, position in batch.rows: + self._require_prefix(allocation, layer, depth, position + 1) + plan = batch.reference_plan + if plan is not None and not plan.dense(): + return sparse_paged_attention( + q, + self.key_cache[:, layer], + self.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + dense_attention=self.attention, + ) if isinstance(batch, DeviceKVMetadata): return self.attention( q, @@ -728,8 +744,6 @@ def _attend_prepared( ) if not batch.rows: return torch.empty_like(q) - for allocation, depth, position in batch.rows: - self._require_prefix(allocation, layer, depth, position + 1) if batch.cu_seqlens_q is not None: return self.attention.prefill( q, diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py index f8be4a75..3c9d977b 100644 --- a/vllm_rlt/engine/preemption.py +++ b/vllm_rlt/engine/preemption.py @@ -43,10 +43,18 @@ def suspend(self, victim, *, offload_kv=True): pipeline = runner.snapshot_pipeline(victim, device) elif pipeline is not None: pipeline = pipeline.to(device) + # Sparse decode: a token past its reference loop owns a reference that + # its remaining loops read. It moves with the rest of the saved state. + references = previous.get("references") if previous is not None else None + if previous is None and e.reference_store is not None: + references = e.model_runner.snapshot_attention_reference(victim, device) + elif references is not None: + references = (references[0], references[1].to(device)) # Build a new record before releasing resources. In particular, a failed # offload must leave the prior resident state available for resumption. snapshot = dict( pipeline=pipeline, + references=references, phase=phase, kv=cache.snapshot(victim.request_id) if offload_kv else None, hidden=None if hidden is None else hidden.to(device).clone(), @@ -85,6 +93,8 @@ def resume(self, request): ) if self.engine.speculative_runner is not None: self.engine.speculative_runner.restore_pipeline(request, state.get("pipeline")) + if state.get("references") is not None: + self.engine.model_runner.restore_attention_reference(request, state["references"]) # Resident copies completed before suspension released their runner slot. # Restored CPU buffers must be visible to every stream before resumption. if kv is not None and cache.device.type == "cuda": diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..0140f73f 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -1,7 +1,10 @@ import logging +import math from dataclasses import dataclass, replace from typing import TYPE_CHECKING, cast +import torch + from vllm_rlt.attention import BackendCapabilities, create_backend from vllm_rlt.config import ( CacheConfig, @@ -20,6 +23,7 @@ from vllm_rlt.profiling import Profiler from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.attention_reference import AttentionReferenceStore from vllm_rlt.worker.executor import Executor from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput @@ -51,6 +55,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + sparse_decode_attention_config=None, stage_config=None, connector: StageConnector | None = None, ): @@ -87,6 +92,7 @@ def __init__( speculative_config and speculative_config.mode == "pipeline" ) self.async_speculative = bool(speculative_config and self.execution_config.async_scheduling) + scratch = 0 if speculative_config is not None: if cache_config.layout != "last_exited": raise ValueError("speculative decoding requires last_exited KV") @@ -118,6 +124,149 @@ def __init__( scratch = 2 * scratch + 8 * scheduler_config.max_num_seqs * ( config.total_ut_steps * width ) + sparse_decode = sparse_decode_attention_config + self.sparse_decode_attention = sparse_decode + self.reference_store = None + if sparse_decode is not None and sparse_decode.enabled: + if getattr(model, "requires_boundary_kv", False): + raise ValueError("sparse decode attention requires a recurrent Ouro-style model") + if cache_config.layout != "last_exited": + raise ValueError("sparse decode attention requires last_exited KV") + # Graph replays run only the fused kernels, which need no host + # work per row; other paths would silently fall back to dense. + graphs = self.execution_config.cuda_graphs + if graphs and ( + not self.async_speculative + or sparse_decode.implementation != "fused" + or sparse_decode.key_retention >= 1.0 + ): + raise ValueError( + "sparse decode attention with CUDA graphs requires async block " + "speculation, the fused implementation and key_retention < 1" + ) + if self.pipeline_speculative and ( + not self.async_speculative + or sparse_decode.implementation != "fused" + or sparse_decode.key_retention >= 1.0 + ): + # Pipeline rows move and restart on the device; only the fused + # kernels read their device-decided modes and slots. + raise ValueError( + "sparse decode attention with token pipelines requires asynchronous " + "scheduling, the fused implementation and key_retention < 1" + ) + if scheduler_config.mode != "refill": + raise ValueError("sparse decode attention requires refill scheduling") + if ( + scheduler_config.enable_preemption + and speculative_config is not None + and not self.pipeline_speculative + ): + # Ordinary decode and token pipelines snapshot their live + # references; block rounds hold per-round leases with none. + raise ValueError( + "sparse decode attention supports preemption only for ordinary " + "decode and token pipelines" + ) + # The index buffer never truncates the algorithm: capacity must + # hold the worst-case retained count over the supported serving + # context, which the caller may bound explicitly instead of paying + # for the model's full positional limit. + context_bound = sparse_decode.max_context_tokens or config.max_position_embeddings + if context_bound > config.max_position_embeddings: + raise ValueError( + f"max_context_tokens {context_bound} exceeds the model context " + f"{config.max_position_embeddings}" + ) + required_keys = math.ceil(context_bound * sparse_decode.key_retention) + capacity = sparse_decode.key_bucket or required_keys + if capacity < required_keys: + raise ValueError( + f"sparse decode attention key_bucket {capacity} cannot hold " + f"{required_keys} retained keys for key_retention " + f"{sparse_decode.key_retention}; raise key_bucket or lower key_retention" + ) + # Pool capacity follows simultaneously live inputs, not requests + # alone. Synchronous block references retire in the same round; + # asynchronous speculation can retain two rounds per request. + slots = scheduler_config.max_num_seqs + if speculative_config is not None: + slots *= speculative_config.num_speculative_tokens + 1 + if self.async_speculative: + slots *= 2 + # Account for the reference pool before admitting any work: it is + # allocated up front, so a budget that cannot fit must fail here + # with an actionable message instead of OOM-ing mid-batch. + reference_bytes = ( + (slots + graphs) + * config.num_hidden_layers + * config.num_attention_heads + * (capacity * 4 + 4 + 4 * 2 * config.head_dim + 8) + ) + if parameter.device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(parameter.device) + if reference_bytes > free_bytes: + raise ValueError( + f"sparse decode attention reference pool needs " + f"{reference_bytes / 1024**3:.1f} GiB but only " + f"{free_bytes / 1024**3:.1f} GiB is free; lower " + f"max_context_tokens, key_retention or num_speculative_tokens" + ) + self.reference_store = AttentionReferenceStore( + slots=slots, + layers=config.num_hidden_layers, + heads=config.num_attention_heads, + head_dim=config.head_dim, + key_bucket=capacity, + # Graphs bake the index pool's address in: allocate it whole. + initial_key_capacity=capacity if graphs else min(capacity, 64), + device=parameter.device, + implementation=sparse_decode.implementation, + padding_slot=graphs, + ) + if parameter.device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(parameter.device) + growth_bytes = self.reference_store.peak_capacity_bytes - self.reference_store.bytes + if growth_bytes > free_bytes: + raise ValueError( + "sparse reference index growth exceeds available device memory" + ) + logger.info( + "sparse decode attention: %d reference slots x %d keys x %d layers x %d heads " + "<= %.2f GiB growth peak (context bound %d, retention %.2f)", + slots, + capacity, + config.num_hidden_layers, + config.num_attention_heads, + self.reference_store.peak_capacity_bytes / 1024**3, + context_bound, + sparse_decode.key_retention, + ) + # Budget the logical maximum and copy overlap, not just the small + # initial allocation; short prompts must not overcommit long ones. + scratch += self.reference_store.peak_capacity_bytes + if graphs: + from vllm_rlt.kernels.sparse_decode import SELECT_PARTS + from vllm_rlt.kernels.triton_attention import NUM_SPLITS + from vllm_rlt.worker.cuda_graph import speculative_capture_sizes + + # One capture workspace serves every bucket of the graphs. + rows = speculative_capture_sizes( + self.execution_config, scheduler_config, speculative_config + )[-1] + width = math.ceil(config.max_position_embeddings / cache_config.block_size) + head_dim = 1 << (config.head_dim - 1).bit_length() + scratch += ( + 4 + * rows + * config.num_attention_heads + * (width * cache_config.block_size + NUM_SPLITS * (2 + head_dim)) + ) + # Select scratch per row and head: radix histograms, part counts, + # offsets and contribution parts, plus state, lse and ticket. + parts = SELECT_PARTS * (4 * 256 + 1 + 2 + 1 + head_dim) * 4 + scratch += rows * config.num_attention_heads * (parts + 5 * 2 * 8 + 8) + if speculative_config is not None or self.reference_store is not None: cache_config = replace( cache_config, memory_reserve_bytes=cache_config.memory_reserve_bytes + scratch ) @@ -167,6 +316,7 @@ def __init__( parameter.device.type == "cuda" and self.execution_config.async_scheduling ), device_context_lengths=self.async_speculative, + sparse_decode=self.reference_store is not None, ), ) num_blocks, self.memory_plan = plan_cache( @@ -217,6 +367,8 @@ def __init__( scheduler_config=scheduler_config, staged_prefill=self.scheduler.staged_prefill, rng=self.rng, + reference_store=self.reference_store if speculative_config is None else None, + sparse_decode=self.sparse_decode_attention, ) self.speculative_runner = ( SpeculativeRunner( @@ -227,6 +379,8 @@ def __init__( rng=self.rng, runtime=self.model_runner.runtime, compute_stream=self.model_runner.core_stream, + reference_store=self.reference_store, + sparse_decode=self.sparse_decode_attention, ) if speculative_config is not None and not self.async_speculative else None @@ -244,7 +398,15 @@ def __init__( copy_stream=self.model_runner.copy_stream, rng=self.rng, runtime=self.model_runner.runtime, + reference_store=self.reference_store, + sparse_decode=self.sparse_decode_attention, ) + if self.reference_store is not None: + # Reference lifecycle device work is ordered on the compute + # stream, after the readers that retire a slot. + self.reference_store.stream = self.model_runner.core_stream + if self.reference_store is not None: + self.reference_store.stream = self.model_runner.core_stream self.executor = Executor(self.model_runner, speculation=self.speculative_runner) self._pending_speculative = [] self._retiring_speculative = {} @@ -354,6 +516,12 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i capacity = len(prompt_token_ids) + params.max_tokens - 1 if capacity > config.max_position_embeddings: raise ValueError("prompt plus decode positions exceed the model context length") + if self.reference_store is not None: + bound = ( + self.sparse_decode_attention.max_context_tokens or config.max_position_embeddings + ) + if capacity > bound: + raise ValueError("request exceeds sparse decode attention context bound") cache = self.cache_manager required = cache.required_blocks(capacity) if required > cache.num_blocks: @@ -400,6 +568,10 @@ def close(self): self.abort_request(rid) if self.async_speculative: self._discard_speculative() + if self.reference_store is not None: + # Abort can enqueue reference metadata retirement after the drain + # above. Teardown must also cover those final writes. + self.model_runner.synchronize() def _discard_speculative(self): """Release speculative owners after the caller has drained both streams.""" diff --git a/vllm_rlt/engine/worker.py b/vllm_rlt/engine/worker.py index 9a2b2805..e04429d8 100644 --- a/vllm_rlt/engine/worker.py +++ b/vllm_rlt/engine/worker.py @@ -194,9 +194,16 @@ def run(self): and not self.channel.poll() ): time.sleep(0.001) + graphs = getattr(getattr(self.engine, "speculative_runner", None), "graphs", None) + replays = dict( + graph_replays=getattr(graphs, "replays", 0), + sparse_graph_replays=getattr(graphs, "sparse_replays", 0), + ) self.engine.close() metrics = self.connector.close() if self.connector is not None else {} - self.send("stopped", used_blocks=self.engine.cache_manager.num_used_blocks, **metrics) + self.send( + "stopped", used_blocks=self.engine.cache_manager.num_used_blocks, **replays, **metrics + ) def worker_main( diff --git a/vllm_rlt/entrypoints/llm.py b/vllm_rlt/entrypoints/llm.py index 55e46104..8c4c83b5 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, + sparse_decode_attention_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, + sparse_decode_attention_config=sparse_decode_attention_config, ) self._next_request_id = 0 diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 883d884b..5544d31a 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -5,7 +5,13 @@ from pathlib import Path from types import SimpleNamespace -from vllm_rlt.config import CacheConfig, ExecutionConfig, ExitConfig, SpeculativeConfig +from vllm_rlt.config import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) from vllm_rlt.profiling import ProfileConfig @@ -48,6 +54,22 @@ def add_runtime_args(parser): parser.add_argument("--exit-trace", help="JSON file with depths_by_request") parser.add_argument("--lookahead-seed", type=int, default=0) parser.add_argument("--async-scheduling", action="store_true") + parser.add_argument( + "--sparse-decode-attention", + action="store_true", + help="Recompute decode attention only for retained reference-loop keys", + ) + parser.add_argument( + "--sparse-decode-key-retention", + type=float, + default=1.0, + help="Retained key fraction per query/head; 1.0 is the full-retention control", + ) + parser.add_argument( + "--sparse-decode-max-context", + type=int, + help="Supported serving context bound for the reference index buffer", + ) parser.add_argument( "--cuda-graphs", action="store_true", @@ -129,6 +151,13 @@ def runtime_configs(args): static_buffers=args.static_buffers, pad_to_power_of_two=args.pad_to_power_of_two, ), + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=args.sparse_decode_attention, + key_retention=args.sparse_decode_key_retention, + max_context_tokens=args.sparse_decode_max_context, + ) + if args.sparse_decode_attention + else None, ) diff --git a/vllm_rlt/kernels/sparse_capture.py b/vllm_rlt/kernels/sparse_capture.py new file mode 100644 index 00000000..e9671f35 --- /dev/null +++ b/vllm_rlt/kernels/sparse_capture.py @@ -0,0 +1,152 @@ +"""Fused paged score pass for the sparse-decode capture step. + +The selector needs a probability for every key of the reference loop, but the +Torch oracle materialises the whole causal row in FP32 with the GQA heads +expanded, which costs tens of MiB per layer and leaves the pass CPU-bound. This +kernel reads the paged cache once and writes one FP32 score per +``(row, head, position)``, so selection, softmax and the retained contribution +stay on the device without an expanded copy. + +It computes the same dot product as the oracle in FP32 (the cache is BF16, and +BF16 -> FP32 is exact); only the reduction order differs. Measured against the +oracle at 8k context: max score delta 1.1e-5 on a magnitude of 65, retained mass +delta 9e-8, full attention output delta 2e-8, and identical selected sets on +511/512 row-head pairs. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _capture_scores_kernel( + q_ptr, + k_ptr, + t_ptr, + l_ptr, + out_ptr, + context, + q_stride_row, + k_stride_block, + DIM: tl.constexpr, + GROUPS: tl.constexpr, + KV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + row = tl.program_id(0) + kv_head = tl.program_id(1) + chunk = tl.program_id(2) + + position = chunk * BLOCK_N + tl.arange(0, BLOCK_N) + row_length = tl.load(l_ptr + row) + mask = position < row_length + page = tl.load(t_ptr + row * PAGES + position // BLOCK_SIZE, mask=mask, other=0) + offset = position % BLOCK_SIZE + dim = tl.arange(0, DIM) + # [BLOCK_N, DIM] + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + for group in tl.static_range(GROUPS): + head = kv_head * GROUPS + group + query = tl.load(q_ptr + row * q_stride_row + head * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store( + out_ptr + (row * (KV * GROUPS) + head) * context + position, + scores, + mask=mask, + ) + + +def fused_capture_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + *, + layer: int, + block_size: int, + block_n: int = 128, + context: int | None = None, + ragged: bool | None = None, +) -> torch.Tensor: + """Scores of every row against its causal keys, straight from the cache. + + ``queries`` is ``[rows, heads, dim]``, ``key_cache`` is either + ``[blocks, layers, block_size, kv_heads, dim]`` or the already sliced + ``[blocks, block_size, kv_heads, dim]``, and ``block_tables`` is + ``[rows, pages]``. Returns ``[rows, heads, max_context]`` FP32 with entries + beyond each row's length left at zero. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused capture needs Triton") + if not queries.is_cuda: + raise ValueError("fused capture is a CUDA path") + + rows, heads, dim = queries.shape + if key_cache.dim() == 5: + layer_cache = key_cache[:, layer] + kv_heads = key_cache.shape[3] + else: + # The cache manager hands the attention layer a per-layer view already. + layer_cache = key_cache + kv_heads = key_cache.shape[2] + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + lengths = context_lengths.to(torch.int64) + # ``context`` and ``ragged`` are host-known for a runner submission. Reading + # them back with ``.item()`` would sync the stream once per layer, which + # RFC #156 section 6.2 forbids, so callers pass them in. + if context is None: + context = int(lengths.max().item()) if lengths.numel() else 0 + if context == 0: + return queries.new_zeros((rows, heads, 0), dtype=torch.float32) + + pages = block_tables.shape[1] + out = torch.zeros(rows, heads, context, device=queries.device, dtype=torch.float32) + grid = (rows, kv_heads, triton.cdiv(context, block_n)) + _capture_scores_kernel[grid]( + queries.contiguous(), + layer_cache, + block_tables, + lengths, + out, + context, + queries.stride(0), + key_cache.stride(0), + DIM=dim, + GROUPS=groups, + KV=kv_heads, + BLOCK_SIZE=block_size, + PAGES=pages, + BLOCK_N=block_n, + ) + # Rows are ragged; the kernel leaves the tail at zero, so mask it to -inf + # before any softmax sees it. + if ragged is None: + ragged = lengths.numel() > 0 and int(lengths.min().item()) != context + if ragged: + positions = torch.arange(context, device=queries.device)[None, None, :] + out = out.masked_fill(positions >= lengths[:, None, None], float("-inf")) + return out diff --git a/vllm_rlt/kernels/sparse_consume.py b/vllm_rlt/kernels/sparse_consume.py new file mode 100644 index 00000000..06b7fc7b --- /dev/null +++ b/vllm_rlt/kernels/sparse_consume.py @@ -0,0 +1,209 @@ +"""Fused gather-and-accumulate for the sparse-decode consume step. + +A consume row only touches its ``k`` retained keys, but the Torch path gathers +them into a dense ``[rows, heads, k, dim]`` tensor twice (K then V) and runs one +cuBLAS call per row. At 32k context with ``k = 3277`` that is 8.2 ms of K gather +plus 8.1 ms of V gather through three advanced indices on a strided +``[blocks, block_size, kv, dim]`` cache, which measures about 45 GB/s. + +These two kernels walk the retained indices in place: the first writes one FP32 +score per ``(row, head, retained)``, the second accumulates ``p * V`` into one +output row. Nothing dense is materialised, and the softmax in between stays in +Torch because ``k`` is small. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _selected_scores_kernel( + q_ptr, + k_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + k_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + block = tl.program_id(2) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + + slot = block * BLOCK_K + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + dim = tl.arange(0, DIM) + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + query = tl.load(q_ptr + (row * heads + head) * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store(out_ptr + (row * heads + head) * retained + slot, scores, mask=mask) + + @triton.jit + def _selected_values_kernel( + p_ptr, + v_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + v_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + dim = tl.arange(0, DIM) + accumulator = tl.zeros([DIM], dtype=tl.float32) + for start in range(0, retained, BLOCK_K): + slot = start + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + weight = tl.load( + p_ptr + (row * heads + head) * retained + slot, mask=mask, other=0.0 + ).to(tl.float32) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + values = tl.load( + v_ptr + + page[:, None] * v_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + accumulator += tl.sum(weight[:, None] * values, axis=0) + tl.store(out_ptr + (row * heads + head) * DIM + dim, accumulator) + + +def _cache_view(cache: torch.Tensor, layer: int) -> tuple[torch.Tensor, int]: + """Return the per-layer cache view and its KV head count.""" + if cache.dim() == 5: + return cache[:, layer], cache.shape[3] + return cache, cache.shape[2] + + +def selected_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``q . K[S]`` for every retained key, without materialising K. + + ``queries`` is ``[rows, heads, dim]``, ``selected`` is ``[rows, heads, k]`` + of absolute positions and the result is ``[rows, heads, k]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, dim = queries.shape + cache, kv_heads = _cache_view(key_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + retained = selected.shape[-1] + out = torch.empty(rows, heads, retained, device=queries.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads, triton.cdiv(retained, block_k)) + _selected_scores_kernel[grid]( + queries.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out + + +def selected_values( + probabilities: torch.Tensor, + value_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``sum_j p[j] * V[S[j]]`` for every row and head. + + ``probabilities`` is ``[rows, heads, k]`` and the result is + ``[rows, heads, dim]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, retained = probabilities.shape + cache, kv_heads = _cache_view(value_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + dim = cache.shape[-1] + out = torch.empty(rows, heads, dim, device=probabilities.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads) + _selected_values_kernel[grid]( + probabilities.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py new file mode 100644 index 00000000..7af6f64f --- /dev/null +++ b/vllm_rlt/kernels/sparse_decode.py @@ -0,0 +1,1499 @@ +"""Fused sparse decode attention for uniform capture or consume submissions. + +Opt-in alternative to the batched path in ``vllm_rlt.attention.sparse``. That +path keeps the scalar oracle's cuBLAS reduction shapes, so every layer launches +per-row products and reads row lengths and retained counts back to the host. +Here each layer launches a fixed number of kernels and keeps lengths and +counts on the device. + +Numerics differ from the oracle on consume rows: scores and FP32 reductions +run in another order, so a key on the selection boundary can differ from the +oracle's. Retained positions are written in ascending order and equal scores +keep the smaller positions, as in the oracle. Capture rows still return the +dense backend output, which is also published as ``O_ref``, so the recurrence +input of the reference loop is unchanged. + +A consume row whose reference was never captured (count 0) yields NaN instead +of raising: raising would need the host readback this path exists to avoid. The +batched and scalar paths keep the explicit error. +""" + +import operator + +import torch +import triton +import triton.language as tl + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DeviceSparseDecodePlan, retained_count +from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention +from vllm_rlt.kernels.triton_attention import NUM_SPLITS + +# Row modes as the kernels read them from a device plan. +CAPTURE_MODE = tl.constexpr(CAPTURE) +CONSUME_MODE = tl.constexpr(CONSUME) + + +@triton.jit +def _order_key(scores): + """Key in [0, 2**32) with the order of FP32 scores (bits of negatives flipped).""" + bits = scores.to(tl.int32, bitcast=True) + return (bits ^ ((bits >> 31) & 0x7FFFFFFF)).to(tl.int64) + 2147483648 + + +@triton.jit +def _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + reference, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Keep the ``count`` largest of ``length`` stored scores and publish ``(S, m, O_ref, C_ref)``. + + Finds the count-th largest score by a radix select over its 32-bit order + key, four bits per pass, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. ``reference`` is ``O_ref``. + """ + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + # After eight passes ``prefix`` is the order key of the count-th largest + # score and ``ties_needed`` how many keys equal to it are still retained. + # Loop-carried values must keep one type on every iteration. + prefix = length.to(tl.int64) * 0 + ties_needed = count + digits = tl.arange(0, 16) + for level in tl.static_range(8): + histogram = tl.zeros((16,), dtype=tl.int32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 4 * level)) == prefix) + digit = ((key >> (28 - 4 * level)) & 15).to(tl.int32) + histogram += tl.sum( + ((digit[:, None] == digits[None, :]) & candidate[:, None]).to(tl.int32), axis=0 + ) + # Highest digit whose keys, with every higher digit, still cover the rest. + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + prefix = prefix * 16 + chosen + low = prefix + selected = SELECTED + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head) * retained_capacity + written = 0 + ties_seen = 0 + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + scores = tl.load(scores_row + positions, mask=live, other=0) + key = _order_key(scores) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + probabilities = tl.where(take, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=take, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=take[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + tl.store(MASS + state, mass) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, contribution, mask=dims < HEAD_DIM) + + +# A copy of triton_attention._paged_attention_kernel that also stores the raw +# scores. Its arithmetic must stay identical to the dense kernel: a uniform +# capture's output is required to equal the dense backend bit for bit (tested). +# The mixed-mode variant (device plans) is compiled separately and matches it to +# fp32 rounding; some architectures (sm_89) round it differently by an ulp. +@triton.jit +def _capture_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + PART_M, + PART_L, + PART_O, + SCORES, + SEGHIST, + RETENTION, + STATE, + TICKETS, + MODES, + top, + score_capacity, + q_batch_stride: tl.constexpr, + q_head_stride: tl.constexpr, + q_dim_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + k_dim_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + v_dim_stride: tl.constexpr, + table_stride: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + NUM_HEADS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, + HISTOGRAM: tl.constexpr, + PARTS: tl.constexpr, + MIXED: 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) + if MIXED: + # Per-row modes: consume rows do no dense work, dense rows keep no scores. + mode = tl.load(MODES + row) + length = tl.where(mode == CONSUME_MODE, 0, length) + # Equal tile-aligned parts of this row's context; trailing parts may be empty. + split_tokens = ((length + SPLITS - 1) // SPLITS + BLOCK_T - 1) // BLOCK_T * BLOCK_T + begin = split * split_tokens + end = tl.minimum(length, begin + split_tokens) + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_batch_stride + head * q_head_stride + dims * q_dim_stride, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + local = tl.zeros((512,), dtype=tl.int32) + for start in range(begin, end, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + valid_tokens = positions < end + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, + mask=valid_tokens, + other=0, + ).to(tl.int64) + # Physical IDs fit in int32, but ID * block_stride can exceed 2**31 + # elements in a large KV pool. Promote BEFORE multiplying, for K and V. + 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, :] * k_dim_stride, + 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")) + # The only addition to the dense kernel: keep the raw scores for selection. + keep = valid_tokens + if MIXED: + keep = keep & (mode == CAPTURE_MODE) + tl.store( + SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + positions, + scores, + mask=keep, + ) + if HISTOGRAM: + # Selection pass 0 of this part: the top eight bits of the order key. + digit = (_order_key(scores) >> 24).to(tl.int32) + local += tl.histogram(tl.where(valid_tokens, digit, 256), 512) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :] * v_dim_stride, + 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 + # An empty part stores (-inf, 0, 0), which the merge weights by zero. + part = (row * NUM_HEADS + head) * SPLITS + split + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + if HISTOGRAM: + rh = row * NUM_HEADS + head + _, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED) + if count > 0: + bins = tl.arange(0, 512) + tl.store(SEGHIST + (rh * PARTS + split) * 4 * 256 + bins, local, mask=bins < 256) + _close_pass(SEGHIST, STATE, TICKETS, rh, count, 0, SPLITS, PARTS) + + +# Parallel selection after the split kernel. Every (row, head) splits its +# context into PARTS tile-aligned parts. Four radix passes over the 32-bit order +# key, eight bits each, store one 256-bin histogram per part. The last part to +# finish a pass (an integer ticket per row and head) sums the histograms in a +# fixed order and publishes the state for the next pass, so the result stays +# deterministic and no histogram is zeroed. After the last pass it also derives +# every part's output offset and tie quota, which keeps the ascending order and +# the smaller-position tie rule. The contribution kernel then gathers V over the +# compact retained list only. +SELECT_PARTS = 16 +# Tiles of the radix passes and the write (scores only) and of the V gather, +# chosen by a sweep on an RTX 5090 at 4K-32K (16 vs 32 parts, 128-1024 scan, +# 32-128 gather, 4 or 8 warps). +_SCAN_BLOCK = 1024 +_SCAN_WARPS = 4 +_GATHER_BLOCK = 128 +# The split kernel's tile (the dense kernel's), so its splits are whole parts. +_SEGMENT_ALIGN = 32 +# The split kernel builds the first pass's histograms while it scores the keys. +_PASS0_IN_SPLIT = True + +# Triton's Python dispatch costs several microseconds per launch, and a fused +# capture makes six launches, so at short contexts in eager mode the dispatch, +# not the GPU, sets its cost. A launch whose specialization was seen before +# reuses that compiled kernel directly. The key follows Triton's own +# specialization (dtype and 16-byte alignment of pointers; integers equal to 1, +# divisible by 16, and their width; constexpr values); tests set _LAUNCH_CHECK +# to assert that it always picks the kernel Triton would. Other Triton versions +# keep its dispatch. +try: + from triton import knobs as _knobs + from triton.runtime.driver import driver as _driver + + _FAST_LAUNCH = triton.__version__.startswith("3.8.") and hasattr( + _knobs.runtime, "launch_enter_hook" + ) +except ImportError: # pragma: no cover - other Triton versions + _FAST_LAUNCH = False +_LAUNCH_CHECK = False + + +def _key_function(args, constexpr): + """Compile the specialization key of one argument layout, as Triton's binder does. + + The layout is the arguments' types; each layout gets its own function. + """ + terms = [] + for i, value in enumerate(args): + if i in constexpr: + terms.append(f"a[{i}]") + elif isinstance(value, torch.Tensor): + terms.append(f"a[{i}].dtype, a[{i}].data_ptr() & 15 == 0") + elif type(value) is int: + terms.append(f"a[{i}] == 1, a[{i}] & 15 == 0, -2147483648 <= a[{i}] < 2147483648") + elif type(value) is not float: + return None + if not terms: + return lambda a: () + body = ", ".join(terms) + return eval(f"lambda a: ({body},)") # noqa: S307 - built from indices only + + +class _Launcher: + """``kernel[grid](*args, **meta)`` that skips Triton's dispatch once compiled.""" + + def __init__(self, kernel): + self.kernel = kernel + self.compiled = {} + self.keys = {} + # Positional parameters declared constexpr specialize on their value. + self.constexpr = frozenset(i for i, p in enumerate(kernel.params) if p.is_constexpr) + + def launch(self, grid, args, meta): + if not _FAST_LAUNCH: + self.kernel[grid](*args, **meta) + return + layout = (tuple(map(type, args)), len(grid)) + keyed = self.keys.get(layout) + if keyed is None: + keyed = self.keys[layout] = _key_function(args, self.constexpr) + if keyed is None: + self.kernel[grid](*args, **meta) + return + device = _driver.active.get_current_device() + key = (device, layout, keyed(args), tuple(meta.items())) + entry = self.compiled.get(key) + if entry is None: + compiled = self.kernel[grid](*args, **meta) + # The launcher takes the bound arguments in the binder's order. + binder = self.kernel.device_caches[device][4] + options = dict( + meta, + debug=self.kernel.debug or _knobs.runtime.debug, + instrumentation_mode=_knobs.compilation.instrumentation_mode, + ) + names = tuple(binder(*args, **options)[0])[len(args) :] + # None when the call passes its meta in the binder's order. + if names == tuple(meta): + order = None + elif len(names) > 1: + order = operator.itemgetter(*names) + else: + order = lambda meta: tuple(meta[name] for name in names) # noqa: E731 + self.compiled[key] = (compiled, order) + return + compiled, order = entry + if _LAUNCH_CHECK: + assert self.kernel.run(*args, grid=grid, warmup=True, **meta) is compiled + values = args + (tuple(meta.values()) if order is None else order(meta)) + stream = _driver.active.get_current_stream(device) + compiled.run( + grid[0], + grid[1] if len(grid) > 1 else 1, + grid[2] if len(grid) > 2 else 1, + stream, + compiled.function, + compiled.packed_metadata, + compiled.launch_metadata(grid, stream, *values), + _knobs.runtime.launch_enter_hook, + _knobs.runtime.launch_exit_hook, + *values, + ) + + +def _current_stream(device): + """The raw current stream; Triton's driver query is much cheaper than torch's.""" + if _FAST_LAUNCH: + return _driver.active.get_current_stream(device.index) + return torch.cuda.current_stream(device).cuda_stream + + +_LAUNCHERS: dict = {} + + +def _launch(kernel, grid, *args, **meta): + # Keyed by identity: hashing a JITFunction recomputes its cache key. + launcher = _LAUNCHERS.get(id(kernel)) + if launcher is None or launcher.kernel is not kernel: + launcher = _LAUNCHERS[id(kernel)] = _Launcher(kernel) + launcher.launch(grid, args, meta) + + +@triton.jit +def _retained(LENGTHS, row, RETENTION, top): + """The row length and its retained count, as the host's ``retained_count``.""" + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + return length, tl.minimum(count.to(tl.int32), top) + + +@triton.jit +def _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED: tl.constexpr): + """``_retained`` for a capture row; a mixed submission's other rows retain nothing.""" + length, count = _retained(LENGTHS, row, RETENTION, top) + if MIXED: + count = tl.where(tl.load(MODES + row) == CAPTURE_MODE, count, 0) + return length, count + + +@triton.jit +def _segment(length, SEGMENTS: tl.constexpr, ALIGN: tl.constexpr): + """This program's part ``[first, last)`` of a row's context.""" + size = tl.cdiv(tl.cdiv(length, SEGMENTS), ALIGN) * ALIGN + first = tl.program_id(2) * size + return first, tl.minimum(first + size, length) + + +@triton.jit +def _choose(histogram, prefix, ties_needed): + """Highest digit whose keys, with every higher digit, still cover the rest.""" + digits = tl.arange(0, 256) + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + return prefix * 256 + chosen.to(tl.int64), ties_needed, chosen + + +@triton.jit +def _last_part(TICKETS, rh, PARTS: tl.constexpr): + """Whether this program is the last of its row and head to finish; resets the ticket. + + Each part stores its results first; the acquire-release ticket then makes + every part's stores visible to the last one, which reads them through L2. + """ + tl.debug_barrier() + last = tl.atomic_add(TICKETS + rh, 1, sem="acq_rel", scope="gpu") == PARTS - 1 + if last: + tl.store(TICKETS + rh, 0) + return last + + +@triton.jit +def _close_pass( + SEGHIST, STATE, TICKETS, rh, count, PASS: tl.constexpr, USED: tl.constexpr, PARTS: tl.constexpr +): + """The last part of a pass publishes the state before the next pass.""" + if _last_part(TICKETS, rh, USED): + digits = tl.arange(0, 256) + parts = (rh * PARTS + tl.arange(0, USED)) * 4 + PASS + histogram = tl.sum( + tl.load(SEGHIST + parts[:, None] * 256 + digits[None, :], cache_modifier=".cg"), axis=0 + ) + if PASS == 0: + prefix = count.to(tl.int64) * 0 + ties_needed = count + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + ties_needed = tl.load(STATE + (rh * 5 + PASS) * 2 + 1).to(tl.int32) + prefix, ties_needed, _ = _choose(histogram, prefix, ties_needed) + tl.store(STATE + (rh * 5 + PASS + 1) * 2, prefix) + tl.store(STATE + (rh * 5 + PASS + 1) * 2 + 1, ties_needed.to(tl.int64)) + + +@triton.jit +def _select_pass_kernel( + PART_M, + PART_L, + PART_O, + OUT, + SCORES, + LENGTHS, + RETENTION, + SLOTS, + COUNTS, + OUTPUT, + VALID, + LSE, + SEGHIST, + STATE, + TICKETS, + GREATER, + OFFSETS, + MODES, + layer, + top, + score_capacity, + out_batch_stride: tl.constexpr, + out_head_stride: tl.constexpr, + out_dim_stride: tl.constexpr, + PASS: tl.constexpr, + MERGE: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, + SPLITS: tl.constexpr, + MIXED: tl.constexpr, +): + """One radix pass of one part; the first launched pass also merges the split parts.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED) + if MERGE: + merging = segment == 0 + marking = True + if MIXED: + # Consume rows get their output from the consume kernels. + mode = tl.load(MODES + row) + merging = merging & (mode != CONSUME_MODE) + marking = mode == CAPTURE_MODE + if merging: + # Identical to triton_attention._paged_attention_merge_kernel. + parts = rh * SPLITS + tl.arange(0, SPLITS) + dims = tl.arange(0, BLOCK_D) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + dense = accumulator / tl.maximum(normalizer, 1.0e-30) + out = OUT + row * out_batch_stride + head * out_head_stride + dims * out_dim_stride + tl.store(out, dense, mask=dims < HEAD_DIM) + slot = tl.load(SLOTS + row) + if head == 0: + if marking: + tl.store(VALID + slot * NUM_LAYERS + layer, 1) + tl.store(LSE + rh, maximum + tl.log(normalizer)) + if count > 0: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + # O_ref is the output as rounded to the query dtype: the recurrence input. + reference = dense.to(OUT.dtype.element_ty).to(tl.float32) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + if count > 0: + if PASS == 0: + prefix = length.to(tl.int64) * 0 + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + first, last = _segment(length, PARTS, ALIGN) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + local = tl.zeros((512,), dtype=tl.int32) + greater = 0 + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 8 * PASS)) == prefix) + digit = ((key >> (24 - 8 * PASS)) & 255).to(tl.int32) + # Bin 256 collects the non-candidates and is never stored. + local += tl.histogram(tl.where(candidate, digit, 256), 512) + if PASS == 3: + # Keys of this part already above the threshold's first 24 bits. + greater += tl.sum((live & ((key >> 8) > prefix)).to(tl.int32), axis=0) + bins = tl.arange(0, 512) + own = (rh * PARTS + segment) * 4 + PASS + tl.store(SEGHIST + own * 256 + bins, local, mask=bins < 256) + if PASS == 3: + tl.store(GREATER + rh * PARTS + segment, greater) + if _last_part(TICKETS, rh, PARTS): + # Fix the threshold, then every part's offset and tie quota. + digits = tl.arange(0, 256) + segments = tl.arange(0, PARTS) + prefix = tl.load(STATE + (rh * 5 + 3) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 3) * 2 + 1).to(tl.int32) + last_pass = tl.load( + SEGHIST + ((rh * PARTS + segments[:, None]) * 4 + 3) * 256 + digits[None, :], + cache_modifier=".cg", + ) + low, ties_needed, chosen = _choose(tl.sum(last_pass, axis=0), prefix, ties_needed) + above = tl.load(GREATER + rh * PARTS + segments, cache_modifier=".cg") + tl.sum( + tl.where(digits[None, :] > chosen, last_pass, 0), axis=1 + ) + equal = tl.sum(tl.where(digits[None, :] == chosen, last_pass, 0), axis=1) + ties_before = tl.cumsum(equal, 0) - equal + taken = above + tl.minimum(tl.maximum(ties_needed - ties_before, 0), equal) + tl.store(OFFSETS + (rh * PARTS + segments) * 2, tl.cumsum(taken, 0) - taken) + tl.store(OFFSETS + (rh * PARTS + segments) * 2 + 1, ties_before) + tl.store(STATE + (rh * 5 + 4) * 2, low) + tl.store(STATE + (rh * 5 + 4) * 2 + 1, ties_needed.to(tl.int64)) + else: + _close_pass(SEGHIST, STATE, TICKETS, rh, count, PASS, PARTS, PARTS) + + +@triton.jit +def _select_write_kernel( + SCORES, + LENGTHS, + RETENTION, + SLOTS, + STATE, + OFFSETS, + SELECTED, + MODES, + layer, + top, + score_capacity, + retained_capacity, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, + MIXED: tl.constexpr, +): + """Write this part's retained positions at its offset, ties up to its quota.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED) + if count > 0: + low = tl.load(STATE + (rh * 5 + 4) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 4) * 2 + 1).to(tl.int32) + written = tl.load(OFFSETS + (rh * PARTS + segment) * 2) + ties_seen = tl.load(OFFSETS + (rh * PARTS + segment) * 2 + 1) + slot = tl.load(SLOTS + row) + selected = ( + SELECTED + + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head).to(tl.int64) * retained_capacity + ) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + first, last = _segment(length, PARTS, ALIGN) + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + + +@triton.jit +def _select_contribution_kernel( + SCORES, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + LSE, + SELECTED, + PART_MASS, + PART_C, + TICKETS, + MASS, + CONTRIBUTION, + MODES, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + MIXED: tl.constexpr, +): + """``(m, C_ref)`` over the retained ranks, in rank chunks summed in a fixed order. + + Gathering V over the compact retained list keeps every lane busy, where a + scan over all positions would leave most of them masked. + """ + row = tl.program_id(0) + head = tl.program_id(1) + part = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED) + if count > 0: + chunk = tl.cdiv(tl.cdiv(count, PARTS), BLOCK_T) * BLOCK_T + first = part * chunk + last = tl.minimum(first + chunk, count) + lse = tl.load(LSE + rh) + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + selected = SELECTED + state.to(tl.int64) * retained_capacity + scores_row = SCORES + rh.to(tl.int64) * score_capacity + table = TABLES + row * table_stride + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + scores = tl.load(scores_row + positions, mask=live, other=0) + probabilities = tl.where(live, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + index = rh * PARTS + part + tl.store(PART_MASS + index, mass) + tl.store(PART_C + index * BLOCK_D + dims, contribution) + if _last_part(TICKETS, rh, PARTS): + parts = rh * PARTS + tl.arange(0, PARTS) + total = tl.sum(tl.load(PART_MASS + parts, cache_modifier=".cg"), axis=0) + summed = tl.sum( + tl.load(PART_C + parts[:, None] * BLOCK_D + dims[None, :], cache_modifier=".cg"), + axis=0, + ) + tl.store(MASS + state, total) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, summed, mask=dims < HEAD_DIM) + + +@triton.jit +def _capture_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + DENSE, + SCORES, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + VALID, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + dense_row_stride: tl.constexpr, + dense_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Select and publish the retained state ``(S, m, O_ref, C_ref)`` of one row and head. + + One program scores every key, finds the ``count``-th largest score by + bisection on its order key, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. It also marks the reference source-ready for this layer, as + ``AttentionReferenceStore.capture`` does, without a separate launch. + """ + row = tl.program_id(0) + head = tl.program_id(1) + slot = tl.load(SLOTS + row) + if head == 0: + tl.store(VALID + slot * NUM_LAYERS + layer, 1) + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + # key_bucket covers the serving context, so this never truncates; it only + # keeps a misconfiguration from writing past the index buffer. + count = tl.minimum(count.to(tl.int32), top) + if count > 0: + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + table = TABLES + row * table_stride + scores_row = SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + maximum = -float("inf") + normalizer = 0.0 + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, 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=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + tl.store(scores_row + positions, scores, mask=live) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + normalizer = normalizer * tl.exp(maximum - next_maximum) + tl.sum( + tl.exp(scores - next_maximum), axis=0 + ) + maximum = next_maximum + lse = maximum + tl.log(normalizer) + # Later passes read scores other threads of this program stored. + tl.debug_barrier() + dense = tl.load( + DENSE + row * dense_row_stride + head * dense_head_stride + dims, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + dense, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride, + v_token_stride, + v_head_stride, + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + HEAD_GROUPS, + PAGE_SIZE, + BLOCK_D, + BLOCK_T, + ) + + +@triton.jit +def _consume_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + SLOTS, + SELECTED, + COUNTS, + PART_M, + PART_L, + PART_O, + MODES, + layer, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + 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, + MIXED: tl.constexpr, +): + """Online-softmax parts of ``C_new`` over one contiguous chunk of retained ranks. + + Each (row, head) splits its ``count`` retained keys into ``SPLITS`` chunks + of whole ``BLOCK_T`` tiles; an empty chunk writes the merge's neutral part. + """ + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + dims = tl.arange(0, BLOCK_D) + part = (row * NUM_HEADS + head) * SPLITS + split + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + reads = tl.load(LENGTHS + row) > 0 + if MIXED: + reads = reads & (tl.load(MODES + row) == CONSUME_MODE) + if reads: + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + count = tl.load(COUNTS + state) + chunk = tl.cdiv(tl.cdiv(count, SPLITS), BLOCK_T) * BLOCK_T + first = split * chunk + last = tl.minimum(first + chunk, count) + selected = SELECTED + state.to(tl.int64) * retained_capacity + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, mask=live, 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=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, 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 + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + + +@triton.jit +def _consume_merge_kernel( + PART_M, + PART_L, + PART_O, + LENGTHS, + SLOTS, + MASS, + OUTPUT, + CONTRIBUTION, + USES, + OUT, + MODES, + layer, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + SPLITS: tl.constexpr, + MIXED: tl.constexpr, +): + """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order. + + Also counts the use, as ``AttentionReferenceStore.consume`` does. + """ + row = tl.program_id(0) + head = tl.program_id(1) + dims = tl.arange(0, BLOCK_D) + out = OUT + row * q_row_stride + head * q_head_stride + dims + # A mixed submission writes only its consume rows here. + consuming = True + if MIXED: + consuming = tl.load(MODES + row) == CONSUME_MODE + if consuming: + slot = tl.load(SLOTS + row) + if head == 0: + uses = USES + slot * NUM_LAYERS + layer + tl.store(uses, tl.load(uses) + 1) + if tl.load(LENGTHS + row) == 0: + tl.store(out, tl.zeros((BLOCK_D,), dtype=tl.float32), mask=dims < HEAD_DIM) + else: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + parts = (row * NUM_HEADS + head) * SPLITS + tl.arange(0, SPLITS) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + reference = tl.load(OUTPUT + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0) + contribution = tl.load( + CONTRIBUTION + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0 + ) + # count 0 (no captured reference) divides 0 by 0: NaN, never a silent value. + mass = tl.load(MASS + state) + corrected = reference - contribution + mass * (accumulator / normalizer) + tl.store(out, corrected, mask=dims < HEAD_DIM) + + +_RETENTION: dict[tuple[torch.device, float], torch.Tensor] = {} + + +def _retention(device, key_retention): + """A float64 device scalar, built once; Triton would pass a Python float as fp32.""" + key = (device, key_retention) + if key not in _RETENTION: + _RETENTION[key] = torch.full((1,), key_retention, dtype=torch.float64, device=device) + return _RETENTION[key] + + +_WORKSPACE: dict = {} + + +def _build_capture_workspace(device, rows, heads, capacity, block_d): + empty = lambda *size, dtype=torch.float32: torch.empty(size, device=device, dtype=dtype) # noqa: E731 + return ( + empty(rows, heads, capacity), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + # Per-part radix histograms, the state before each pass (and the final + # threshold), per-part keys above the threshold and output offsets, the + # merged lse, the last-part tickets (zero between calls), and the + # contribution parts. + empty(rows, heads, SELECT_PARTS, 4, 256, dtype=torch.int32), + empty(rows, heads, 5, 2, dtype=torch.int64), + empty(rows, heads, SELECT_PARTS, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS, 2, dtype=torch.int32), + empty(rows, heads), + torch.zeros(rows, heads, device=device, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS), + empty(rows, heads, SELECT_PARTS, block_d), + ) + + +def _capture_workspace(device, rows, heads, capacity, block_d): + """Capture scratch reused across layers. + + Captures issued on one stream run in order, so each layer may overwrite + the previous layer's scratch; one entry per stream keeps memory bounded. + """ + stream = _current_stream(device) + shape = (rows, heads, capacity, block_d) + entry = _WORKSPACE.get((device, stream)) + if entry is None or entry[0] != shape: + entry = (shape, *_build_capture_workspace(device, rows, heads, capacity, block_d)) + _WORKSPACE[(device, stream)] = entry + return entry[1:] + + +_CONSUME_WORKSPACE: dict = {} + + +def _consume_workspace(device, rows, heads, block_d): + """Consume split parts, grown per stream and reused across layers.""" + stream = _current_stream(device) + entry = _CONSUME_WORKSPACE.get((device, stream)) + if ( + entry is None + or entry[0].shape[0] < rows + or entry[0].shape[1] != heads + or entry[2].shape[-1] != block_d + ): + empty = lambda *size: torch.empty(size, device=device, dtype=torch.float32) # noqa: E731 + entry = ( + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + ) + _CONSUME_WORKSPACE[(device, stream)] = entry + return (None, *entry) + + +def supports(plan, q, dense_attention) -> bool: + """Whether the fused path covers this submission; otherwise use the batched path.""" + if isinstance(plan, DeviceSparseDecodePlan): + # Dense and capture rows run the split kernel; it needs the dense backend. + return ( + q.is_cuda + and q.shape[-1] <= 256 + and plan.key_retention < 1.0 + and dense_attention is not None + ) + modes = set(plan.modes) + return ( + q.is_cuda + and q.shape[-1] <= 256 + and plan.key_retention < 1.0 + and (modes == {CONSUME} or (modes == {CAPTURE} and dense_attention is not None)) + ) + + +def capture_workspace(device, rows, heads, capacity, dim): + """Capture scratch for at most ``rows`` rows, owned by the caller. + + A CUDA graph bakes the scratch addresses in, so its owner keeps one + workspace alive for every bucket instead of the shape-keyed cache. + """ + return _build_capture_workspace(device, rows, heads, capacity, triton.next_power_of_2(dim)) + + +def fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention, + workspace=None, +): + """Capture or consume one submission without host readbacks. + + A host plan has one mode for every row. A device plan + (``DeviceSparseDecodePlan``) carries each row's mode on the device: the + capture kernels then compute dense and capture rows, the consume kernels + consume rows, and each writes only its own rows of the output. + + The kernels also write the store's ``valid``/``uses`` marks, so the host + only checks the plan's liveness instead of launching the bookkeeping. + ``workspace`` comes from ``capture_workspace`` and must cover the rows + and table width of the submission; consume uses only its split parts. + """ + store = plan.store + store._require_live(plan) + rows, heads, dim = q.shape + page = key_cache.shape[1] + groups = heads // key_cache.shape[2] + block_d = triton.next_power_of_2(dim) + scale = dim**-0.5 + for cache in (key_cache, value_cache): + if cache.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous head dimension") + if q.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous query head dimension") + selected = store.selected + if not selected.is_contiguous(): + raise ValueError("fused sparse attention derives index offsets from a contiguous pool") + mixed = isinstance(plan, DeviceSparseDecodePlan) + # Uniform plans never read the mode pointer; any device tensor stands in. + modes = plan.modes if mixed else context_lengths + if mixed or plan.modes[0] == CAPTURE: + # With the Triton backend, the dense split kernel also keeps the scores + # (and the first radix histograms), the dense merge runs inside the + # first selection pass, and selection runs in parallel parts. Other + # backends keep their call plus one serial selection kernel. + # Engines pass their backend adapter, which wraps the Triton kernel; + # callers may also pass the Triton function itself. + kernel = getattr(dense_attention, "kernel", dense_attention) + split = isinstance(kernel, TritonPagedAttention) or kernel is triton_paged_attention + if mixed and not split: + raise ValueError("a device sparse plan needs the Triton dense attention backend") + if split: + output = torch.empty_like(q) + else: + output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + if rows == 0: + return output + # The round's page tables bound every row's context on the host. + capacity = block_tables.shape[1] * page + top = min(retained_count(capacity, plan.key_retention), store.key_bucket or capacity) + store.ensure_key_capacity(top) + selected = store.selected + if workspace is None: + workspace = _capture_workspace(q.device, rows, heads, capacity, block_d) + elif workspace[0].shape[1:] != (heads, capacity) or workspace[0].shape[0] < rows: + raise ValueError("capture workspace does not cover the submission") + ( + scores, + part_m, + part_l, + part_o, + seghist, + state, + greater, + offsets, + lse, + tickets, + part_mass, + part_c, + ) = workspace if workspace[0].shape[0] == rows else (buffer[:rows] for buffer in workspace) + if split: + folded = _PASS0_IN_SPLIT + retention = _retention(q.device, plan.key_retention) + _launch( + _capture_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + part_m, + part_l, + part_o, + scores, + seghist, + retention, + state, + tickets, + modes, + top, + capacity, + *q.stride(), + *key_cache.stride(), + *value_cache.stride(), + block_tables.stride(0), + HEAD_DIM=dim, + HEAD_GROUPS=groups, + NUM_HEADS=heads, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + HISTOGRAM=folded, + PARTS=SELECT_PARTS, + MIXED=mixed, + num_warps=4, + ) + first_pass = 1 if folded else 0 + # A device plan whose step has no capture loop still merges its dense + # rows in the first pass; it launches no further selection kernels. + selecting = not mixed or plan.may_capture + for level in range(first_pass, 4 if selecting else first_pass + 1): + _launch( + _select_pass_kernel, + (rows, heads, SELECT_PARTS), + part_m, + part_l, + part_o, + output, + scores, + context_lengths, + retention, + plan.slot_ids, + store.counts, + store.output, + store.valid, + lse, + seghist, + state, + tickets, + greater, + offsets, + modes, + layer, + top, + capacity, + *output.stride(), + PASS=level, + MERGE=level == first_pass, + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + SPLITS=NUM_SPLITS, + MIXED=mixed, + num_warps=_SCAN_WARPS, + ) + if selecting: + _launch( + _select_write_kernel, + (rows, heads, SELECT_PARTS), + scores, + context_lengths, + retention, + plan.slot_ids, + state, + offsets, + selected, + modes, + layer, + top, + capacity, + selected.shape[-1], + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + MIXED=mixed, + num_warps=_SCAN_WARPS, + ) + _launch( + _select_contribution_kernel, + (rows, heads, SELECT_PARTS), + scores, + value_cache, + block_tables, + context_lengths, + retention, + plan.slot_ids, + lse, + selected, + part_mass, + part_c, + tickets, + store.mass, + store.contribution, + modes, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + BLOCK_D=block_d, + BLOCK_T=_GATHER_BLOCK, + PARTS=SELECT_PARTS, + MIXED=mixed, + num_warps=4, + ) + if not mixed: + return output + else: + # Serial selection after another backend's dense call (uniform plans). + _launch( + _capture_kernel, + (rows, heads), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + _retention(q.device, plan.key_retention), + plan.slot_ids, + output, + scores, + selected, + store.counts, + store.mass, + store.output, + store.contribution, + store.valid, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + 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), + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=64, + num_warps=4, + ) + return output + if not mixed: + output = torch.empty_like(q) + if rows: + if workspace is None: + workspace = _consume_workspace(q.device, rows, heads, block_d) + elif workspace[1].shape[0] < rows or workspace[1].shape[1] != heads: + raise ValueError("consume workspace does not cover the submission") + part_m, part_l, part_o = ( + workspace[1:4] if workspace[1].shape[0] == rows else (b[:rows] for b in workspace[1:4]) + ) + _launch( + _consume_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan.slot_ids, + selected, + store.counts, + part_m, + part_l, + part_o, + modes, + layer, + block_tables.stride(0), + selected.shape[-1], + 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), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + MIXED=mixed, + num_warps=4, + ) + _launch( + _consume_merge_kernel, + (rows, heads), + part_m, + part_l, + part_o, + context_lengths, + plan.slot_ids, + store.mass, + store.output, + store.contribution, + store.uses, + output, + modes, + layer, + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + SPLITS=NUM_SPLITS, + MIXED=mixed, + num_warps=4, + ) + return output diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index 47f58ad1..3814f8f0 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -69,6 +69,7 @@ def __init__( model_config=None, seed=0, speculative_config=None, + sparse_decode_attention_config=None, ): self.speculative_config = speculative_config model, revision, self.model_config = prepare_model( @@ -138,6 +139,11 @@ def __init__( else self.exit_config ), speculative_config=speculative_config if role == "decode" else None, + # Reference state is captured and consumed while decoding; + # prefill rows are always dense, so only D owns a store. + sparse_decode_attention_config=( + sparse_decode_attention_config if role == "decode" else None + ), execution_config=execution, attention_backend=attention_backend, ) diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 6b708df2..95857fa5 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -8,6 +8,7 @@ import numpy as np import torch +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, REFERENCE_DEPTH, reference_mode from vllm_rlt.core.kv_cache_manager import _Allocation from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import Operation, Request @@ -36,6 +37,9 @@ class DevicePipelineState: hidden: torch.Tensor candidates: torch.Tensor active: torch.Tensor + # Sparse decode: a permutation of the request's reference slots, moved with + # the tokens, so each row keeps the slot its loop-1 capture wrote. + slots: torch.Tensor | None = None @dataclass @@ -50,6 +54,12 @@ class SpeculativeRequestState: event: torch.cuda.Event | None = None # Row in the runner's device slot table holding (token, position). slot: int | None = None + # Bumped per request instance so reference store bindings reject a reused + # request ID across generations. + generation: int = 0 + # Token pipelines own their reference slots for the request's lifetime: a + # row captures in one submission and consumes in the next. + reference: object | None = None class SpeculativeRoundBuffer: @@ -92,6 +102,10 @@ class SpeculativeRoundTicket: stats: SpeculativeStats | None = None cached: ModelRunnerOutput | None = None retired: bool = False + # Reference leases of this round. Released on retire, after the final + # reader completed, so no later submission can observe or overwrite them. + references: tuple = () + reference_store: object | None = None @property def has_result(self): @@ -148,6 +162,9 @@ def retire(self, *, discard=False): # collect() already waited for copy_done; normal delivery adds no wait. self.buffer.keepalive.clear() self.buffer.leased = False + if self.reference_store is not None and self.references: + self.reference_store.release(self.references) + self.references = () self.retired = True @@ -158,6 +175,11 @@ class AsyncSpeculativeRunner(SpeculativeRunner): reservation upper bound and must never be used as the actual write position. """ + # Sparse decode attention lifecycle; test doubles that skip __init__ + # inherit the disabled defaults. + reference_store = None + sparse_decode = None + def __init__( self, model, @@ -170,6 +192,8 @@ def __init__( copy_stream, rng: RngRegistry | None = None, runtime=None, + reference_store=None, + sparse_decode=None, ): super().__init__( model, @@ -180,9 +204,24 @@ def __init__( runtime=runtime, compute_stream=compute_stream, ) + # Sparse decode attention lifecycle: the runner owns the per-token + # reference pool (RFC #156 section 5) and drives retirement from the + # GPU acceptance counts without a host readback. + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self._generations: dict[str, int] = {} sizes = speculative_capture_sizes(execution, scheduler, config) self.graphs = ( - DeviceRecurrentGraphs(model, cache, execution, False, capture_sizes=sizes) + DeviceRecurrentGraphs( + model, + cache, + execution, + False, + capture_sizes=sizes, + reference_store=reference_store, + key_retention=None if sparse_decode is None else sparse_decode.key_retention, + device_plans=config.mode == "pipeline", + ) if execution.cuda_graphs else None ) @@ -222,15 +261,24 @@ def _state(self, request, buffer): host = torch.tensor([request.input_token_id, request.position], pin_memory=True) values = host.to(self.device, non_blocking=True) buffer.keepalive.extend((host, values)) + generation = self._generations.get(rid, 0) + 1 + self._generations[rid] = generation if self.config.mode == "pipeline": - state = SpeculativeRequestState(request, allocation, values[0], values[1]) + state = SpeculativeRequestState( + request, allocation, values[0], values[1], generation=generation + ) else: if not self.free_slots: self._grow_slots() slot = self.free_slots.pop() self.slot_values[slot].copy_(values) state = SpeculativeRequestState( - request, allocation, self.slot_values[slot, 0], self.slot_values[slot, 1], slot=slot + request, + allocation, + self.slot_values[slot, 0], + self.slot_values[slot, 1], + slot=slot, + generation=generation, ) if self.config.mode == "pipeline": width = min( @@ -249,6 +297,9 @@ def _state(self, request, buffer): torch.zeros(width, dtype=torch.long, device=self.device), torch.arange(width, device=self.device) == 0, ) + if self.reference_store is not None: + state.reference = self.reference_store.reserve(rid, 0, width, generation) + state.pipeline.slots = state.reference.slot_ids.clone() self.states[rid] = state return state @@ -264,7 +315,9 @@ def _grow_slots(self): for state in self.states.values(): state.token, state.position = new[state.slot, 0], new[state.slot, 1] - def _execute_recurrent_step(self, hidden, positions, tables, depth, allocations, active=None): + def _execute_recurrent_step( + self, hidden, positions, tables, depth, allocations, active=None, reference_plan=None + ): if isinstance(depth, torch.Tensor): table = tables.gather(1, depth[:, None, None].expand(-1, 1, tables.shape[2])).squeeze(1) else: @@ -280,6 +333,7 @@ def _execute_recurrent_step(self, hidden, positions, tables, depth, allocations, (positions + 1) if active is None else torch.where(active, positions + 1, 0) ).int(), write_mask=active, + reference_plan=reference_plan, ) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden, metadata=(metadata,), graphs=self.graphs @@ -365,9 +419,47 @@ def add(name, values): values = host.to(self.device, non_blocking=True) buffer.keepalive.extend((host, values)) views = {name: values[start : start + size] for name, (start, size) in layout.items()} - return SimpleNamespace(n=n, width=width, rows=len(request_rows), offsets=offsets, **views) + return SimpleNamespace( + n=n, width=width, rows=len(request_rows), offsets=offsets, counts=counts, **views + ) - def _draft(self, plan, tables, allocations): + def _reference_plan(self, leases, rows, depth): + """Build the sparse decode plan of one recurrent step. + + ``rows`` holds one ``(lease_index, row)`` pair per step row in step + order. Dense steps return ``None`` so their metadata stays untouched. + """ + if self.reference_store is None or leases is None: + return None + mode = reference_mode(depth) + if mode == DENSE: + return None + owners = [(leases[lease_index], row) for lease_index, row in rows] + return self.reference_store.build_plan( + owners, [mode] * len(owners), key_retention=self.sparse_decode.key_retention + ) + + def _pipeline_reference_plan(self, states, slots, depth, active, offset): + """Device modes of one pipeline step: each row's loop decides its mode.""" + if slots is None: + return None + if any(state.reference is None for state in states): + raise RuntimeError("pipeline reference slots were discarded with their buffers") + capture = active & (depth == REFERENCE_DEPTH) + consume = active & (depth > REFERENCE_DEPTH) + modes = torch.where(capture, CAPTURE, torch.where(consume, CONSUME, DENSE)).to(torch.int8) + # Rows sit at multiples of draft_loops, so whether any row can be in its + # capture loop at this offset is known on the host. + depths = range(0, self.config.target_loops, self.config.draft_loops) + return self.reference_store.build_device_plan( + [state.reference for state in states], + slots, + modes, + key_retention=self.sparse_decode.key_retention, + may_capture=REFERENCE_DEPTH - offset in depths, + ) + + def _draft(self, plan, tables, allocations, leases=None): from vllm_rlt.kernels.routing import gather_kernel slot_values = self.slot_values.new_empty((plan.n, 2)) @@ -394,9 +486,18 @@ def _draft(self, plan, tables, allocations): positions = starts[active] + offset row_tables = tables[active] hidden = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=row_tokens)).hidden + # A drafting offset covers one verification row per active request: + # the token at column `offset`, which the verifier later reuses. + rows = range(plan.n) if active is None else [int(index) for index in active] for depth in range(self.config.draft_loops): + reference = self._reference_plan(leases, [(i, offset) for i in rows], depth) hidden = self._execute_recurrent_step( - hidden, positions, row_tables, depth, allocations + hidden, + positions, + row_tables, + depth, + allocations, + reference_plan=reference, ) hidden_blocks.append(hidden) if drafting_count: @@ -406,15 +507,31 @@ def _draft(self, plan, tables, allocations): tokens[getattr(plan, f"drafting{offset}"), offset + 1] = proposed return starts, tokens, hidden_blocks - def _verify(self, plan, starts, hidden_blocks, tables, allocations): + def _verify(self, plan, starts, hidden_blocks, tables, allocations, leases=None): positions = starts[plan.request_rows] + plan.row_offsets row_tables = tables[plan.request_rows] hidden = torch.cat(hidden_blocks)[plan.hidden_rows] + # Verification rows are request-major, matching the round plan. + verification_rows = [ + (i, j) for i, count in enumerate(plan.counts) for j in range(int(count)) + ] for depth in range(self.config.draft_loops, self.config.target_loops): - hidden = self._execute_recurrent_step(hidden, positions, row_tables, depth, allocations) + reference = self._reference_plan(leases, verification_rows, depth) + if reference is not None and depth == self.config.draft_loops: + # One device-owned absolute position per reference slot; the + # verification rows cover every slot of the round exactly once. + self.reference_store.bind_positions(reference, positions) + hidden = self._execute_recurrent_step( + hidden, + positions, + row_tables, + depth, + allocations, + reference_plan=reference, + ) return self._coda(hidden).argmax(-1) - def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer): + def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer, leases=None): from vllm_rlt.kernels.routing import scatter_kernel count = len(target) @@ -426,6 +543,10 @@ def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer) drafts = torch.arange(plan.width - 1, device=self.device)[None, :] < plan.drafts[:, None] matches = (tokens[:, 1:] == targets[:, :-1]) & drafts accepted = matches.long().cumprod(1).sum(1) + if leases is not None and self.reference_store is not None: + # GPU acceptance drives retirement: rejected candidates and their + # dependent descendants lose their references on the device. + self.reference_store.accept(leases, accepted) # All target IDs through the first rejection equal the accepted # candidates followed by the correction (or final bonus). next_tokens = targets.gather(1, accepted[:, None]).squeeze(1) @@ -462,14 +583,40 @@ def submit(self, batch): rounds = tuple(s.submitted for s in states) if any(s.submitted > s.delivered for s in states): self.submitted_before_collect += 1 - tables.copy_(host_tables, non_blocking=True) - if self.config.mode == "pipeline": - count = self._pipeline_step(scheduled_requests, states, tables, allocations, buffer) - else: - plan = self._round_plan(scheduled_requests, states, buffer) - starts, tokens, hidden_blocks = self._draft(plan, tables, allocations) - target = self._verify(plan, starts, hidden_blocks, tables, allocations) - count = self._accept_and_update_state(plan, states, starts, tokens, target, buffer) + # One reference slot per verification row: the round's live tokens. + # Outstanding rounds keep disjoint slots until their tickets retire. + leases = None + if self.reference_store is not None and self.config.mode != "pipeline": + leases = tuple( + self.reference_store.reserve( + scheduled_request.request.request_id, + index, + scheduled_request.token_count, + state.generation, + ) + for index, (scheduled_request, state) in enumerate( + zip(scheduled_requests, states) + ) + ) + try: + tables.copy_(host_tables, non_blocking=True) + if self.config.mode == "pipeline": + count = self._pipeline_step( + scheduled_requests, states, tables, allocations, buffer + ) + else: + plan = self._round_plan(scheduled_requests, states, buffer) + starts, tokens, hidden_blocks = self._draft(plan, tables, allocations, leases) + target = self._verify(plan, starts, hidden_blocks, tables, allocations, leases) + count = self._accept_and_update_state( + plan, states, starts, tokens, target, buffer, leases + ) + except BaseException: + # The engine drains both streams before discarding, so releasing + # here cannot race a later submission for these slots. + if leases is not None: + self.reference_store.release(leases) + raise buffer.compute_done = torch.cuda.Event() buffer.compute_done.record(self.compute_stream) for state in states: @@ -497,6 +644,8 @@ def submit(self, batch): count, mode=self.config.mode, stats=self.stats, + references=leases or (), + reference_store=self.reference_store, ) def _pipeline_step(self, spans, states, tables, allocations, buffer): @@ -518,10 +667,19 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer): prelude = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=tokens)).hidden hidden = torch.where((depths == 0)[:, None], prelude, retained) row_tables = torch.cat([tables[i : i + 1].expand(w, -1, -1) for i, w in enumerate(widths)]) + slots = None + if self.reference_store is not None: + slots = torch.cat([state.pipeline.slots[:w] for state, w in zip(states, widths)]) for offset in range(self.config.draft_loops): depth = torch.where(active, depths + offset, 0) hidden = self._execute_recurrent_step( - hidden, positions, row_tables, depth, allocations, active + hidden, + positions, + row_tables, + depth, + allocations, + active, + reference_plan=self._pipeline_reference_plan(states, slots, depth, active, offset), ) choices = self._coda(hidden).argmax(-1) start = 0 @@ -546,6 +704,10 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer): source = (ids + emitted.long()).clamp_max(len(ids) - 1) remaining = valid.long().sum() - emitted.long() surviving = (ids < remaining) & ~reject + if p.slots is not None: + # A permutation, unlike `source`: a freed root slot moves to the + # tail, where only inactive rows or newborns can take it. + p.slots.copy_(p.slots.gather(0, (ids + emitted.long()) % len(ids))) p.tokens.copy_(p.tokens.gather(0, source)) p.depths.copy_(new_depths.gather(0, source)) p.hidden.copy_(next_hidden[source]) @@ -598,6 +760,17 @@ def snapshot_pipeline(self, request, device): (pipeline.tokens, pipeline.depths, pipeline.candidates, pipeline.active.long()) ).tolist() tokens, depths, candidates, active = values + references = None + if self.reference_store is not None: + if state.reference is None: + raise RuntimeError("sparse pipeline lost its reference slots") + # Restoration puts the k live tokens back in rows 0..k-1; keep the + # slot of each row by recording it relative to the lease. + if not all(active[: sum(active)]): + raise RuntimeError("live pipeline rows must form a prefix") + lease_rows = {slot: row for row, slot in enumerate(state.reference.slots)} + permutation = tuple(lease_rows[slot] for slot in pipeline.slots.tolist()) + references = (self.reference_store.snapshot(state.reference, device), permutation) return PipelineState( request, position, @@ -612,6 +785,7 @@ def snapshot_pipeline(self, request, device): for i, token in enumerate(tokens) if active[i] ], + references=references, ) def restore_pipeline(self, request, snapshot): @@ -619,6 +793,8 @@ def restore_pipeline(self, request, snapshot): return if snapshot.request is not request or request.request_id in self.states: raise RuntimeError("invalid pipeline restoration generation") + if (self.reference_store is not None) != (snapshot.references is not None): + raise RuntimeError("pipeline snapshot and sparse decode configuration disagree") if snapshot.pending or snapshot.position != request.position: raise RuntimeError("pipeline snapshot disagrees with committed request position") width = min( @@ -652,12 +828,27 @@ def restore_pipeline(self, request, snapshot): values[2], torch.arange(width, device=self.device) < len(snapshot.tokens), ) + generation = self._generations.get(request.request_id, 0) + 1 + self._generations[request.request_id] = generation + reference = None + if snapshot.references is not None: + saved, permutation = snapshot.references + if len(saved) != width or sorted(permutation) != list(range(width)): + raise RuntimeError("pipeline reference snapshot does not match the window") + # Fresh slots and generation; each row reads the slot it read + # before suspension. + reference = self.reference_store.restore(saved, request.request_id, 0, generation) + pipeline.slots = reference.slot_ids[ + torch.tensor(permutation, device=self.device) + ].clone() self.states[request.request_id] = SpeculativeRequestState( request, self.cache._get_allocation(request.request_id), torch.tensor(request.input_token_id, device=self.device), torch.tensor(snapshot.position, device=self.device), pipeline, + generation=generation, + reference=reference, ) def request_ready(self, request): @@ -674,6 +865,10 @@ def release(self, request): del self.states[request.request_id] if state.slot is not None: self.free_slots.append(state.slot) + if state.reference is not None: + # The last reader completed with the event above. + with torch.cuda.stream(self.compute_stream): + self.reference_store.release(state.reference) def discard_buffers(self): """Reset round storage after both streams drain and all tickets retire. @@ -685,6 +880,16 @@ def discard_buffers(self): for buffer in self.buffers: buffer.keepalive.clear() buffer.leased = False + # Test doubles borrow this method without the runner __init__. + reference_store = getattr(self, "reference_store", None) + if reference_store is not None: + # Both streams have drained and every ticket retired; any leftover + # binding can only come from a failed submission. + reference_store.discard() + for state in getattr(self, "states", {}).values(): + # Discard freed every binding; a surviving pipeline must not + # release or read these slots again. + state.reference = None def synchronize(self): self.compute_stream.synchronize() diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py new file mode 100644 index 00000000..cadc311c --- /dev/null +++ b/vllm_rlt/worker/attention_reference.py @@ -0,0 +1,681 @@ +"""Per-token attention reference state for sparse decode attention. + +A sparse decode loop reuses the reference loop's attention output and only +recomputes the contribution of the retained keys. The retained state is +per token, per layer and per head, so it cannot be stored per request. + +This module owns that state. The store keeps one slot per live token, tracks +which request and generation owns it, and publishes the slot views that the +attention kernels borrow. Nothing here reads device state back to the host: +ownership, retirement and invalidation are expressed as device operations so +that an asynchronous round can retire its slots without a host readback. + +Slot lifecycle +-------------- +``empty`` -> ``source-ready`` -> ``reusable`` -> ``retired`` + +A slot becomes source-ready when its capture round finished writing the +reference, and reusable once a later loop consumed it. Retirement is ordered +after the last reader of the slot, so a later submission cannot observe a slot +that an earlier round still reads. + +Slot identity +------------- +A slot is identified by its request generation, its row inside the request and +an incarnation counter. The incarnation changes whenever the token of that row +changes, including a corrected token that reuses the same position after a +rejection. A plan built for a superseded incarnation is stale and rejected. +""" + +from contextlib import nullcontext +from dataclasses import dataclass + +import torch + +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + DeviceSparseDecodePlan, + SparseDecodePlan, +) + +__all__ = [ + "ReferenceLease", + "ReferenceSnapshot", + "AttentionReferenceStore", + "DENSE", + "CAPTURE", + "CONSUME", +] + + +@dataclass(frozen=True) +class ReferenceLease: + """One request's reference allocation for one speculative submission. + + ``slots`` maps the request's verification row (token position) to its + reference slot. The device tensors mirror them so that acceptance-driven + retirement stays on the device. + """ + + request_id: str + generation: int + request_index: int + slots: tuple[int, ...] + incarnations: tuple[int, ...] + slot_ids: torch.Tensor + row_index: torch.Tensor + + def __len__(self) -> int: + return len(self.slots) + + +@dataclass(frozen=True) +class ReferenceSnapshot: + """Owned copy of a lease's slots, for suspension across preemption. + + Every per-slot field a later reader or the retirement path looks at is + copied, so a restored slot is indistinguishable from the original one + except for its identity (slot, generation and incarnation are new). + ``host_counts`` keeps the published retained counts of each row. + """ + + selected: torch.Tensor + counts: torch.Tensor + mass: torch.Tensor + output: torch.Tensor + contribution: torch.Tensor + valid: torch.Tensor + uses: torch.Tensor + position: torch.Tensor + accepted: torch.Tensor + host_counts: tuple[dict[int, int], ...] + + def __len__(self) -> int: + return self.counts.shape[0] + + def to(self, device) -> "ReferenceSnapshot": + tensors = { + name: getattr(self, name).to(device) + for name in ( + "selected", + "counts", + "mass", + "output", + "contribution", + "valid", + "uses", + "position", + "accepted", + ) + } + return ReferenceSnapshot(**tensors, host_counts=self.host_counts) + + +class AttentionReferenceStore: + """Device-resident reference state for the live tokens of one stage. + + The store is sized by the engine, which knows how many tokens can be live + at once. It never allocates per request: a speculative round holds one slot + per verification row, and several rounds of the same request may be + outstanding at the same time, so the capacity must cover every outstanding + round. + """ + + def __init__( + self, + *, + slots, + layers, + heads, + head_dim, + key_bucket, + device, + stream=None, + initial_key_capacity=None, + implementation="batched", + padding_slot=False, + ): + if slots < 1 or layers < 1 or heads < 1 or head_dim < 1: + raise ValueError("reference store dimensions must be positive") + if key_bucket is not None and key_bucket < 1: + raise ValueError("key_bucket must be positive when set") + limit = key_bucket or 1 + initial = limit if initial_key_capacity is None else initial_key_capacity + if not 1 <= initial <= limit: + raise ValueError("initial key capacity must be within the reference capacity") + # Conservatively include all growth generations: an asynchronous + # stream may not have finished copying earlier buffers at the next + # growth. The engine budgets this before allocating its KV cache. + capacity = initial + self._peak_selected_capacity = capacity + while capacity < limit: + grown = min(limit, capacity * 2) + self._peak_selected_capacity += grown + capacity = grown + self.slots = slots + self.layers = layers + self.heads = heads + self.head_dim = head_dim + self.key_bucket = key_bucket + self.device = torch.device(device) + self.stream = stream + # Which CUDA attention path consumes this store's plans. + self.implementation = implementation + # CUDA graph replays pad a submission to its bucket. Padding rows read + # and write this extra slot, which is never bound to a token. + self.padding_slot = slots if padding_slot else None + slots += bool(padding_slot) + shape = (slots, layers, heads) + self.selected = torch.zeros((*shape, initial), dtype=torch.int32, device=device) + self.counts = torch.zeros(shape, dtype=torch.int32, device=device) + self.mass = torch.zeros(shape, dtype=torch.float32, device=device) + self.output = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + self.contribution = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + # A reference is usable per layer while its slot is live. + self.valid = torch.zeros((slots, layers), dtype=torch.bool, device=device) + # How many later loops consumed the reference of each slot and layer. + self.uses = torch.zeros((slots, layers), dtype=torch.int32, device=device) + # Identity of the token occupying each slot. + self.position = torch.zeros(slots, dtype=torch.int64, device=device) + self.incarnation = torch.zeros(slots, dtype=torch.int64, device=device) + self.accepted = torch.zeros(slots, dtype=torch.int32, device=device) + self.free = list(range(self.slots)) + self._bindings: dict[int, tuple[str, int, int]] = {} + self._incarnations: dict[tuple[str, int], int] = {} + # Retained counts are host-known: the capture decides them before it + # submits. Keeping them here lets a reader size its gather without a + # device readback, which RFC #156 section 6.2 forbids. + self._counts_host: dict[tuple[int, int], int] = {} + + @property + def bytes(self) -> int: + """Payload bytes held by the store, excluding allocator overhead.""" + return sum( + tensor.numel() * tensor.element_size() + for tensor in ( + self.selected, + self.counts, + self.mass, + self.output, + self.contribution, + self.valid, + self.uses, + self.position, + self.incarnation, + self.accepted, + ) + ) + + @property + def peak_capacity_bytes(self) -> int: + """Conservative payload bound including all in-flight growth generations.""" + index_bytes = ( + self.selected.shape[0] * self.layers * self.heads * self._peak_selected_capacity * 4 + ) + return self.bytes - self.selected.nbytes + index_bytes + + def ensure_key_capacity(self, count: int) -> None: + """Grow shared index storage on the attention owner stream. + + Existing plans borrow the store rather than a fixed index pointer. + Prior reads/writes and the copy execute on the same owner stream; + record_stream protects the old allocation until that copy completes. + The logical key bucket and retained count never shrink with allocation. + """ + limit = self.key_bucket or 1 + if count > limit: + raise ValueError(f"reference capacity {limit} cannot hold {count} retained keys") + current = self.selected.shape[-1] + if count <= current: + return + if self.device.type == "cuda" and self.stream is not None: + if torch.cuda.current_stream(self.device) != self.stream: + raise RuntimeError("reference growth must run on the attention owner stream") + capacity = current + while capacity < count: + capacity = min(limit, capacity * 2) + old = self.selected + grown = torch.zeros((*old.shape[:-1], capacity), dtype=old.dtype, device=old.device) + grown[..., :current].copy_(old) + if old.device.type == "cuda": + old.record_stream(torch.cuda.current_stream(old.device)) + self.selected = grown + + @property + def outstanding(self) -> int: + """Number of slots currently bound to live tokens.""" + return len(self._bindings) + + def _index(self, values) -> torch.Tensor: + """Host-known slot or row indices on the device, without a blocking copy. + + A pageable host-to-device copy waits for the stream to drain, which + serializes the host with every in-flight round; a pinned source lets + the copy queue behind the work already submitted. + """ + host = torch.tensor(values, dtype=torch.int64, pin_memory=self.device.type == "cuda") + return host.to(self.device, non_blocking=True) + + def reserve(self, request_id: str, request_index: int, rows: int, generation: int): + """Take one slot per verification row of a request's submission.""" + if rows < 1: + raise ValueError("a reference reservation needs at least one row") + if len(self.free) < rows: + raise RuntimeError("no retired attention reference slots") + slots = tuple(self.free.pop() for _ in range(rows)) + incarnations = [] + for row, slot in enumerate(slots): + count = self._incarnations.get((request_id, row), 0) + 1 + self._incarnations[(request_id, row)] = count + self._bindings[slot] = (request_id, generation, count) + incarnations.append(count) + slot_ids = self._index(slots) + # Writing host integers element by element would block on each copy. + self.incarnation.index_copy_(0, slot_ids, self._index(incarnations)) + return ReferenceLease( + request_id=request_id, + generation=generation, + request_index=request_index, + slots=slots, + incarnations=tuple(incarnations), + slot_ids=slot_ids, + row_index=torch.arange(rows, dtype=torch.int64, device=self.device), + ) + + def release(self, leases) -> None: + """Return slots to the pool once their last reader completed. + + The caller orders this after the final reader/writer, e.g. by releasing + from a ticket that already waited for the round's completion event. The + cleared payload never leaks into a later incarnation. + """ + if isinstance(leases, ReferenceLease): + leases = (leases,) + cleared = [] + for lease in leases: + for slot in lease.slots: + binding = self._bindings.get(slot) + if binding is None or binding[:2] != (lease.request_id, lease.generation): + raise RuntimeError("attention reference slot released by a different owner") + del self._bindings[slot] + cleared.append(slot) + if cleared: + self.free.extend(cleared) + self._forget_counts(cleared) + if self.stream is None: + self._clear(self._index(cleared)) + return + # A ticket retires from the host's current stream. An index built + # there would be freed, and reusable by that stream, before the + # owner stream's index_fill_ reads it; build it on the owner stream. + with torch.cuda.stream(self.stream): + self._clear(self._index(cleared)) + + def discard(self) -> None: + """Free every binding without device work; the caller has drained.""" + self._bindings.clear() + self._counts_host.clear() + self.free = list(range(self.slots)) + + def _owner(self): + if self.stream is None: + return nullcontext() + return torch.cuda.stream(self.stream) + + def snapshot(self, lease: ReferenceLease, device) -> ReferenceSnapshot: + """Copy a live lease's slots into storage the snapshot owns. + + Suspension calls this after the runner drained, before it releases the + lease. The copy runs on the owner stream and completes before return, + so the released slots can be reused at once. + """ + for slot, incarnation in zip(lease.slots, lease.incarnations): + if self._bindings.get(slot) != (lease.request_id, lease.generation, incarnation): + raise RuntimeError("cannot snapshot a retired attention reference") + with self._owner(): + ids = lease.slot_ids + snapshot = ReferenceSnapshot( + selected=self.selected.index_select(0, ids).to(device), + counts=self.counts.index_select(0, ids).to(device), + mass=self.mass.index_select(0, ids).to(device), + output=self.output.index_select(0, ids).to(device), + contribution=self.contribution.index_select(0, ids).to(device), + valid=self.valid.index_select(0, ids).to(device), + uses=self.uses.index_select(0, ids).to(device), + position=self.position.index_select(0, ids).to(device), + accepted=self.accepted.index_select(0, ids).to(device), + host_counts=tuple( + { + layer: self._counts_host[(slot, layer)] + for layer in range(self.layers) + if (slot, layer) in self._counts_host + } + for slot in lease.slots + ), + ) + if self.stream is not None: + self.stream.synchronize() + return snapshot + + def restore( + self, snapshot: ReferenceSnapshot, request_id: str, request_index: int, generation: int + ) -> ReferenceLease: + """Bind fresh slots for a resumed token and copy its snapshot back.""" + with self._owner(): + # Reserve here too: its slot ids are copied on the current stream, + # which must be the stream that indexes with them below. + lease = self.reserve(request_id, request_index, len(snapshot), generation) + self.ensure_key_capacity(snapshot.selected.shape[-1]) + ids = lease.slot_ids + width = snapshot.selected.shape[-1] + # Only the snapshot's prefix is meaningful; a later capture + # overwrites what it uses, exactly as for a never-suspended slot. + selected = self.selected[ids] + selected[..., :width] = snapshot.selected.to(self.device) + self.selected.index_copy_(0, ids, selected) + for name in ( + "counts", + "mass", + "output", + "contribution", + "valid", + "uses", + "position", + "accepted", + ): + getattr(self, name).index_copy_(0, ids, getattr(snapshot, name).to(self.device)) + if self.stream is not None: + self.stream.synchronize() + for slot, counts in zip(lease.slots, snapshot.host_counts): + for layer, count in counts.items(): + self._counts_host[(slot, layer)] = count + return lease + + def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan: + """Publish the plan of one submission. + + ``owners`` holds one ``(lease, row)`` pair per submission row, in the + row order of the recurrent step; ``modes`` holds the row mode of every + submission row. Slots and incarnations are re-read from the live + bindings, so a plan never aliases a superseded token. + """ + if len(owners) != len(modes): + raise ValueError("a reference plan needs one mode per row") + slots, incarnations = [], [] + for owner, mode in zip(owners, modes): + if owner is None: + if mode != DENSE: + raise ValueError("a sparse row needs a reference owner") + slots.append(-1) + incarnations.append(0) + continue + lease, row = owner + slot = lease.slots[row] + binding = self._bindings.get(slot) + if binding is None or binding != ( + lease.request_id, + lease.generation, + lease.incarnations[row], + ): + raise RuntimeError("attention reference slot released before its plan") + slots.append(slot) + incarnations.append(binding[2]) + return SparseDecodePlan( + modes=tuple(modes), + slots=tuple(slots), + incarnations=tuple(incarnations), + key_retention=key_retention, + key_bucket=self.key_bucket, + slot_ids=self._index(slots), + store=self, + ) + + def build_device_plan( + self, leases, slot_ids, modes, *, key_retention: float, may_capture: bool = True + ): + """Publish a plan whose rows name their slots and modes on the device. + + Every slot in ``slot_ids`` must belong to one of ``leases``; the host + checks those leases live before each use, as for a host plan. + """ + if slot_ids.shape != modes.shape or slot_ids.dim() != 1: + raise ValueError("a device reference plan needs one slot and mode per row") + plan = DeviceSparseDecodePlan( + modes=modes, + slot_ids=slot_ids, + leases=tuple(leases), + key_retention=key_retention, + key_bucket=self.key_bucket, + store=self, + may_capture=may_capture, + ) + self._require_live(plan) + return plan + + def _require_live(self, plan) -> None: + """Reject a plan whose slots were recycled or superseded.""" + if plan.store is not self: + raise RuntimeError("attention reference plan belongs to a different store") + if isinstance(plan, DeviceSparseDecodePlan): + for lease in plan.leases: + for slot, incarnation in zip(lease.slots, lease.incarnations): + binding = self._bindings.get(slot) + if binding != (lease.request_id, lease.generation, incarnation): + raise RuntimeError("stale attention reference plan for a retired token") + return + for slot, incarnation, mode in zip(plan.slots, plan.incarnations, plan.modes): + if slot == -1 and mode == DENSE: + continue + binding = self._bindings.get(slot) + if binding is None or binding[2] != incarnation: + raise RuntimeError("stale attention reference plan for a retired token") + + def note_count(self, slot: int, layer: int, count: int) -> None: + """Record the retained count of one capture row on the host. + + The count is decided by the capture, so it never needs a device + readback; ``host_count`` is the reader-side counterpart. + """ + self._counts_host[(slot, layer)] = count + + def note_counts(self, slot: int, count: int) -> None: + """Record one capture row's retained count for every layer on the host. + + The count follows the row's causal length alone, so a caller that + knows the length can publish it before any layer runs. + """ + for layer in range(self.layers): + self._counts_host[(slot, layer)] = count + + def _forget_counts(self, slots) -> None: + """Drop the host counts of host-known slots, so none sizes a later read.""" + for slot in slots: + for layer in range(self.layers): + self._counts_host.pop((int(slot), layer), None) + + def host_count(self, slot: int, layer: int) -> int | None: + """Retained keys of ``(slot, layer)``, or ``None`` before a capture.""" + return self._counts_host.get((slot, layer)) + + def _row_slots(self, plan: SparseDecodePlan, rows: list[int]) -> torch.Tensor: + """Slot ids of ``rows`` without a host-to-device copy of a Python list. + + ``plan.slot_ids[rows]`` costs ~2.4 ms per call because the list becomes + a device tensor first; grouped rows are contiguous, so derive the index + on the device instead. + """ + if len(rows) == len(plan.slots): + return plan.slot_ids + if rows[0] + len(rows) == rows[-1] + 1: + index = torch.arange(rows[0], rows[0] + len(rows), device=self.device) + return plan.slot_ids.index_select(0, index) + return self._index([plan.slots[row] for row in rows]) + + def capture(self, plan: SparseDecodePlan, layer: int) -> None: + """Mark the captured rows source-ready for the given layer. + + The attention kernel writes the retained state of a capture row and + then calls this on the same stream, so a later reader observes a + complete payload. + """ + self._require_live(plan) + rows = plan.groups().get(CAPTURE) + if not rows: + return + # Assign a tensor, not the Python scalar ``True``: a scalar store into + # a bool tensor through advanced indexing copies it from the host. + slots = self._row_slots(plan, rows) + self.valid[slots, layer] = torch.ones(slots.numel(), dtype=torch.bool, device=self.device) + + def consume(self, plan: SparseDecodePlan, layer: int) -> None: + """Record that the given layer consumed the reference rows.""" + self._require_live(plan) + rows = plan.groups().get(CONSUME) + if not rows: + return + self.uses[self._row_slots(plan, rows), layer] += 1 + + def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> None: + """Record the device-owned absolute positions of the plan's rows.""" + self._require_live(plan) + if positions.shape != plan.slot_ids.shape: + raise ValueError("reference positions need one entry per plan row") + rows = [row for row, slot in enumerate(plan.slots) if slot >= 0] + if len(rows) < len(plan.slots): + positions = positions[self._index(rows)] + self.position[self._row_slots(plan, rows)] = positions.to(torch.int64) + + def write_reference( + self, + plan: SparseDecodePlan, + layer: int, + row: int, + *, + selected, + counts, + mass, + output, + contribution, + ) -> None: + """Publish the retained state of one capture row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = selected.shape[-1] + self.ensure_key_capacity(count) + self.selected[slot, layer, :, :count] = selected.to(torch.int32) + self.counts[slot, layer] = counts.to(torch.int32) + self.note_count(slot, layer, int(count)) + self.mass[slot, layer] = mass + self.output[slot, layer] = output + self.contribution[slot, layer] = contribution + + def read_reference(self, plan: SparseDecodePlan, layer: int, row: int): + """Borrow the retained state of one consume row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = self.host_count(slot, layer) + if count is None: + # A fused capture decides its count on the device. + count = int(self.counts[slot, layer].max()) + if not count: + raise RuntimeError("a consume row has no captured reference") + return ( + self.selected[slot, layer, :, :count].long(), + self.mass[slot, layer], + self.output[slot, layer], + self.contribution[slot, layer], + ) + + def accept(self, leases, accepted: torch.Tensor) -> None: + """Invalidate the references of rejected rows on the device. + + ``accepted`` is the number of accepted candidates of every request of + the submission. A verification row of a request survives exactly when + its row index is at most the accepted count: the root row is always + kept, and a rejected candidate invalidates its dependent descendants. + The comparison stays on the device so that acceptance never requires a + host readback. + + Retirement is memory-bounded: it only clears the per-row metadata that + gates readers (``valid``, ``counts``, ``mass``), never the retained + index buffer. An indexed read-modify-write of ``selected`` would + materialise the whole bucket — at 24 layers, 16 heads and 65,536 keys + that is tens of GiB for one batch — while a zero ``count`` already makes + the payload unreadable, and a later capture overwrites the used prefix. + """ + if isinstance(leases, ReferenceLease): + leases = (leases,) + leases = tuple(leases) + if not leases: + return + slot_ids = torch.cat([lease.slot_ids for lease in leases]) + rows = torch.cat([lease.row_index for lease in leases]) + request = torch.cat( + [torch.full_like(lease.row_index, lease.request_index) for lease in leases] + ) + limit = accepted.to(device=self.device, dtype=rows.dtype).index_select(0, request) + self.accepted[slot_ids] = limit.to(torch.int32) + # A boolean-mask gather of the rejected slots would read its size back + # to the host. Masked writes over the round's own rows stay on device; + # they touch only these rows' per-layer metadata, never the index bucket. + keep = rows <= limit + self.valid[slot_ids] = self.valid[slot_ids] & keep[:, None] + self.counts[slot_ids] = torch.where(keep[:, None, None], self.counts[slot_ids], 0) + self.mass[slot_ids] = torch.where(keep[:, None, None], self.mass[slot_ids], 0) + # Which rows were rejected is known only on the device, so drop the host + # counts of every row of these leases: a later reader then sizes from + # the device counts, which are zero for rejected rows, and a stale host + # count can never size a gather. + self._forget_counts(slot for lease in leases for slot in lease.slots) + + def _clear(self, slots: torch.Tensor) -> None: + """Retire the metadata of the given slots, ordered on the owner stream. + + Only the tensors that gate readers are cleared. The retained index, + output and contribution buffers stay allocated and are overwritten by + the next capture's used prefix, so retiring a slot never writes the + whole bucket (which would be tens of GiB per batch at full retention). + """ + + def zero(): + # index_fill_ passes the scalar to the kernel; assigning a Python + # scalar through an index tensor copies it from the host and blocks. + for tensor in ( + self.valid, + self.uses, + self.accepted, + self.incarnation, + self.position, + self.mass, + self.counts, + ): + tensor.index_fill_(0, slots, 0) + + if self.stream is not None: + with torch.cuda.stream(self.stream): + zero() + else: + zero() + + def invalidate(self, slots) -> None: + """Invalidate the references of the given slots without freeing them.""" + if len(slots) == 0: + return + if isinstance(slots, torch.Tensor): + # Callers hand host lists; a device tensor costs one readback here. + self._forget_counts(slots.tolist()) + slots = slots.to(device=self.device, dtype=torch.int64) + else: + self._forget_counts(slots) + slots = self._index(list(slots)) + for tensor in ( + self.valid, + self.mass, + self.output, + self.contribution, + self.selected, + self.counts, + ): + tensor.index_fill_(0, slots, 0) diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index 10149516..87ae0ec6 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -8,6 +8,13 @@ import torch +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + DeviceSparseDecodePlan, + SparseDecodePlan, +) from vllm_rlt.config import cuda_graph_rows from vllm_rlt.worker.metadata import DeviceKVMetadata @@ -102,8 +109,10 @@ def _capture(hidden, stream, pool, run, *, startup=False): class _DeviceCache: """Capture only GPU cache operations, never Python allocation/written state.""" - def __init__(self, cache): + def __init__(self, cache, workspace=None): self.cache = cache + # Sparse capture scratch shared by every bucket; replays run in order. + self.workspace = workspace def _write_prepared(self, layer, batch, k, v): cache = self.cache @@ -121,6 +130,21 @@ def _write_prepared(self, layer, batch, k, v): def _attend_prepared(self, layer, batch, q): cache = self.cache + plan = getattr(batch, "reference_plan", None) + if plan is not None: + from vllm_rlt.kernels.sparse_decode import fused_paged_attention + + return fused_paged_attention( + q, + cache.key_cache[:, layer], + cache.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + dense_attention=cache.attention, + workspace=self.workspace, + ) if batch.cu_seqlens_q is not None: return cache.attention.prefill( q, @@ -219,10 +243,17 @@ def __init__(self, model, cache, execution, compute_gate): self.last_event = None self.captures = self.replays = self.fallbacks = 0 self.warned_cache_full = False + self.proxy = _DeviceCache(cache) def select(self, hidden, batch): count = _query_count(batch) packed = batch.cu_seqlens_q is not None + plan = getattr(batch, "reference_plan", None) + if plan is not None and not plan.dense(): + # These entries carry no reference plan; replaying one would + # silently run dense attention. + self.fallbacks += 1 + return PreparedGraph(batch) if len(batch.position_ids) != count: batch = replace( batch, @@ -287,13 +318,12 @@ def prepare_inputs(self, call, hidden): stream.wait_event(self.last_event) _copy_graph_inputs(entry, hidden, batch) if entry.graph is None: - proxy = _DeviceCache(cache) entry.graph, entry.output = _capture( entry.hidden, self.capture_stream, self.pool, lambda: self.model.recurrent_prepared( - entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate + entry.hidden, entry.metadata, self.proxy, compute_gate=self.compute_gate ), ) self.captures += 1 @@ -323,16 +353,117 @@ def execute(self, call, hidden): return output +# Graph keys of entries that replay a device plan (token pipelines); not row +# modes. A step in which no row can capture records no selection kernels. +DEVICE_PLAN = "device" +DEVICE_PLAN_NO_CAPTURE = "device-no-capture" + + class DeviceRecurrentGraphs(RecurrentGraphs): """Core graphs whose actual positions and KV frontier remain on device. Allocation identity is checked on every call, outside capture. Entries retain only tensor metadata, never request allocations or CPU written-position state. + + With a fused sparse reference store, every bucket is captured once per row + mode. A sparse entry owns a static plan whose slot ids are refilled before + each replay; the live plan's liveness is checked on the host at that point, + since replay runs no Python. With ``device_plans`` (token pipelines) the + sparse entries instead replay a device plan, whose per-row modes are + refilled alongside the slot ids. """ - def __init__(self, model, cache, execution, compute_gate, capture_sizes=None): + def __init__( + self, + model, + cache, + execution, + compute_gate, + capture_sizes=None, + reference_store=None, + key_retention=None, + device_plans=False, + ): super().__init__(model, cache, execution, compute_gate) self.capture_sizes = capture_sizes or graph_capture_sizes(execution) + self.reference_store = reference_store + self.key_retention = key_retention + self.sparse_replays = 0 + self.modes = (DENSE,) + if reference_store is not None: + from vllm_rlt.kernels.sparse_decode import capture_workspace + + if reference_store.padding_slot is None: + raise ValueError("sparse CUDA graphs need a reference store padding slot") + if reference_store.selected.shape[-1] != (reference_store.key_bucket or 1): + # A replay must never see a grown index pool at a new address. + raise ValueError("sparse CUDA graphs need the reference index pool at full size") + self.modes = ( + (DENSE, DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE) + if device_plans + else (DENSE, CAPTURE, CONSUME) + ) + config = model.config + width = math.ceil(config.max_position_embeddings / cache.block_size) + self.proxy = _DeviceCache( + cache, + capture_workspace( + cache.device, + self.capture_sizes[-1], + config.num_attention_heads, + width * cache.block_size, + config.head_dim, + ), + ) + + @staticmethod + def _key(size, mode): + # Dense entries keep their bucket-size keys; sparse ones add the mode. + return size if mode == DENSE else (size, mode) + + def _entry(self, hidden, batch, size, mode): + entry = _graph_entry( + hidden, + batch, + self.cache, + self.model.config.max_position_embeddings, + count=size, + bucketed=True, + ) + entry.mode = mode + if mode in (DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE): + store = self.reference_store + entry.slot_ids = torch.full( + (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64 + ) + # Padded rows stay dense on the never-bound padding slot. + entry.modes = torch.full((size,), DENSE, device=self.cache.device, dtype=torch.int8) + entry.metadata.reference_plan = DeviceSparseDecodePlan( + modes=entry.modes, + slot_ids=entry.slot_ids, + leases=(), + key_retention=self.key_retention, + key_bucket=store.key_bucket, + store=store, + may_capture=mode == DEVICE_PLAN, + ) + elif mode != DENSE: + store = self.reference_store + entry.slot_ids = torch.full( + (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64 + ) + # Host liveness is checked against the live plan before replay, + # so the captured plan names no host slots. + entry.metadata.reference_plan = SparseDecodePlan( + modes=(mode,) * size, + slots=(), + incarnations=(), + key_retention=self.key_retention, + key_bucket=store.key_bucket, + slot_ids=entry.slot_ids, + store=store, + ) + return entry @torch.inference_mode() def prepare(self): @@ -347,50 +478,77 @@ def prepare(self): block_tables=torch.empty((0, 1), device=self.cache.device, dtype=torch.int32), cu_seqlens_q=None, ) - proxy = _DeviceCache(self.cache) - for size in self.capture_sizes: - entry = _graph_entry( - hidden, - batch, - self.cache, - self.model.config.max_position_embeddings, - count=size, - bucketed=True, - ) - _copy_graph_inputs(entry, hidden, batch, count=0) - entry.graph, entry.output = _capture( - entry.hidden, - self.capture_stream, - self.pool, - lambda: self.model.recurrent_prepared( - entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate - ), - startup=True, - ) - self.entries[size] = entry - self.captures += 1 + for mode in self.modes: + for size in self.capture_sizes: + entry = self._entry(hidden, batch, size, mode) + _copy_graph_inputs(entry, hidden, batch, count=0) + entry.graph, entry.output = _capture( + entry.hidden, + self.capture_stream, + self.pool, + lambda entry=entry: self.model.recurrent_prepared( + entry.hidden, entry.metadata, self.proxy, compute_gate=self.compute_gate + ), + startup=True, + ) + self.entries[self._key(size, mode)] = entry + self.captures += 1 + + def _mode(self, batch): + plan = getattr(batch, "reference_plan", None) + if plan is None or plan.dense(): + return DENSE + if self.reference_store is None or plan.store is not self.reference_store: + return None + if isinstance(plan, DeviceSparseDecodePlan): + if DEVICE_PLAN not in self.modes or plan.key_retention != self.key_retention: + return None + return DEVICE_PLAN if plan.may_capture else DEVICE_PLAN_NO_CAPTURE + modes = set(plan.modes) + if len(modes) != 1 or plan.key_retention != self.key_retention: + return None + return plan.modes[0] def select(self, hidden, batch): self.cache._require_live_batch(batch) count = _query_count(batch) - if not count or count > self.capture_sizes[-1]: + mode = self._mode(batch) + if not count or count > self.capture_sizes[-1] or mode is None: self.fallbacks += 1 return PreparedGraph(batch) capacity = _bucket_size(self.capture_sizes, count) - entry = self.entries.get(capacity) + key = self._key(capacity, mode) + entry = self.entries.get(key) if entry is None: - entry = _graph_entry( - hidden, - batch, - self.cache, - self.model.config.max_position_embeddings, - count=capacity, - bucketed=True, - ) + entry = self._entry(hidden, batch, capacity, mode) entry.graph = None - self.entries[capacity] = entry + self.entries[key] = entry return PreparedGraph(batch, entry) + @torch.inference_mode() + def prepare_inputs(self, call, hidden): + entry = call.entry + if entry is not None and entry.mode != DENSE: + plan = call.metadata.reference_plan + self.reference_store._require_live(plan) + count = _query_count(call.metadata) + stream = torch.cuda.current_stream(self.cache.device) + if self.last_event is not None: + stream.wait_event(self.last_event) + entry.slot_ids[count:].fill_(self.reference_store.padding_slot) + entry.slot_ids[:count].copy_(plan.slot_ids[:count]) + if entry.mode in (DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE): + entry.modes[count:].fill_(DENSE) + entry.modes[:count].copy_(plan.modes[:count]) + super().prepare_inputs(call, hidden) + + @torch.inference_mode() + def execute(self, call, hidden): + sparse = call.entry is not None and call.entry.mode != DENSE + output = super().execute(call, hidden) + self.sparse_replays += sparse + return output + class CodaGraphs: """Select and stage coda inputs before replay; commit KV after submission.""" diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index 8f2c020b..03880166 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -39,6 +39,7 @@ class KVMetadata: writable: bool cu_seqlens_q: torch.Tensor | None = None max_seqlen_q: int = 1 + reference_plan: object | None = None @dataclass(frozen=True, eq=False) @@ -62,6 +63,10 @@ class DeviceKVMetadata: context_lengths: torch.Tensor # Inactive pipeline slots must not overwrite any live KV address. write_mask: torch.Tensor | None = None + # Sparse decode attention plan of this submission: one row mode and + # reference slot per row. Host-known; the attention layer borrows it and + # the runner-owned store keeps the referenced state alive. + reference_plan: object | None = None # The device-owned path currently uses independent single-query rows. # These layout fields also let recurrent graphs stage the same contract. cu_seqlens_q: torch.Tensor | None = None @@ -100,6 +105,7 @@ def prepare_device( block_tables, context_lengths, write_mask=None, + reference_plan=None, ) -> DeviceKVMetadata: """Borrow stream-ordered metadata without reading device positions on CPU. @@ -115,7 +121,10 @@ def prepare_device( block_tables=block_tables, context_lengths=context_lengths, write_mask=write_mask, + reference_plan=reference_plan, ) + if reference_plan is not None and len(reference_plan.modes) != len(position_ids): + raise ValueError("reference plan needs one row mode per query row") if write_mask is not None and ( write_mask.shape != position_ids.shape or write_mask.dtype != torch.bool diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..41a084b1 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -7,6 +7,7 @@ import torch +from vllm_rlt.attention.sparse import CAPTURE, DENSE, reference_mode, retained_count from vllm_rlt.config import ExecutionConfig, ExitConfig, SchedulerConfig from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import ExecutionBatch, Operation, Request @@ -160,9 +161,15 @@ def __init__( scheduler_config=None, staged_prefill=None, rng: RngRegistry | None = None, + reference_store=None, + sparse_decode=None, ): self.model = model.eval() self.cache_manager = cache_manager + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self.attention_references = {} + self.reference_generation = 0 self.metadata = cache_manager.metadata_builder self.runtime = ModelRuntime(self.model, cache_manager) parameter = next(model.parameters()) @@ -1016,6 +1023,8 @@ def _prepare_forward(self, prepared): return replace(prepared, forwards=(PreparedForward(self.runtime.prepare(forward)),)) hidden = self._load_hidden(batch, prepared, workspace, size) if batch.operation == Operation.CODA: + for request in requests: + self._release_attention_reference(request.request_id) forward = ForwardBatch( Operation.CODA, hidden=hidden[: batch.num_rows] if kv_metadata else hidden, @@ -1027,6 +1036,8 @@ def _prepare_forward(self, prepared): kv = prepared.kv if kv is not None: metadata = (kv,) + metadata + if self.reference_store is not None: + metadata = tuple(self._reference_metadata(batch, kv) for kv in metadata) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden[: batch.num_rows] @@ -1040,6 +1051,66 @@ def _prepare_forward(self, prepared): prepared, kv=metadata[0], forwards=(PreparedForward(self.runtime.prepare(forward)),) ) + def _release_attention_reference(self, request_id): + state = self.attention_references.pop(request_id, None) + if state is not None: + # release() creates the index tensor used by the metadata clears. + # Allocate it on core too, so its allocator lifetime follows the + # clear kernels even when CODA or teardown calls from another stream. + with ( + torch.cuda.stream(self.core_stream) + if self.core_stream is not None + else nullcontext() + ): + self.reference_store.release(state[2]) + + def _reference_metadata(self, batch, kv): + """Borrow token-owned references; prompt and padded rows remain dense. + + Preparation runs on core, after the request's preceding event. CODA + retires the reference after its final recurrent read; request teardown + waits for that same event before releasing an unfinished token. + """ + requests = [span.request for span in batch.spans for _ in range(span.token_count)] + owners, modes, captured = [], [], [] + for request, (_, depth, position) in zip(requests, kv.rows): + mode = reference_mode(depth) if position >= len(request.prompt_token_ids) else DENSE + state = self.attention_references.get(request.request_id) + if state is not None and (state[0] is not request or state[1] != position): + self._release_attention_reference(request.request_id) + state = None + if mode == CAPTURE: + if state is None: + self.reference_generation += 1 + lease = self.reference_store.reserve( + request.request_id, 0, 1, self.reference_generation + ) + state = (request, position, lease) + self.attention_references[request.request_id] = state + if state is None: + mode = DENSE # No reference loop for this token: stay dense. + elif mode == CAPTURE: + captured.append((state[2].slots[0], position + 1)) + owners.append(None if mode == DENSE else (state[2], 0)) + modes.append(mode) + if all(mode == DENSE for mode in modes): + return kv + padding = len(kv.position_ids) - len(owners) + owners.extend([None] * padding) + modes.extend([DENSE] * padding) + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode.key_retention + ) + self.reference_store.bind_positions(plan, kv.position_ids) + # A capture row's retained count follows its causal length, which the + # host knows here. Publishing it lets a later batched consume size its + # gather without reading a fused capture's device count back. + for slot, length in captured: + self.reference_store.note_counts( + slot, retained_count(length, self.sparse_decode.key_retention) + ) + return replace(kv, reference_plan=plan) + def _load_hidden(self, batch, prepared, workspace, size): if prepared is not None and prepared.batch_inputs is not None: return prepared.batch_inputs.hidden @@ -1269,6 +1340,8 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu ) except Exception: self._synchronize_submission(stream, prepared) + for span in batch.spans: + self._release_attention_reference(span.request.request_id) # A failed preparation may advance a pool without enqueueing a batch. # Retire all prior readers before any pool can wrap after that gap. self.submissions.drain() @@ -1364,6 +1437,7 @@ def _free_runner_slot(self, request_id): event = self.events.pop(request_id, None) if event is not None: event.synchronize() + self._release_attention_reference(request_id) self.prompt_token_owners.pop(request_id, None) if self.async_state is not None: self.async_state.release(request_id) @@ -1372,6 +1446,35 @@ def _free_runner_slot(self, request_id): if slot is not None: self.free_state_slots.append(slot) + def snapshot_attention_reference(self, request, device): + """Copy the live reference of a token suspended between its loops. + + Call before ``suspend``. A token suspended after its reference loop + must consume that same reference when it resumes; restarting it dense + would be safe but would change its output. + """ + state = self.attention_references.get(request.request_id) + if state is None: + return None + if state[0] is not request: + raise RuntimeError("attention reference belongs to another request generation") + event = self.events.get(request.request_id) + if event is not None: + event.synchronize() + return state[1], self.reference_store.snapshot(state[2], device) + + def restore_attention_reference(self, request, saved): + if saved is None: + return + if request.request_id in self.attention_references: + raise RuntimeError("cannot restore over a live attention reference") + position, snapshot = saved + self.reference_generation += 1 + lease = self.reference_store.restore( + snapshot, request.request_id, 0, self.reference_generation + ) + self.attention_references[request.request_id] = (request, position, lease) + def suspend(self, request_id): """Free the request's runner resources but keep its sampling RNG. diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 85f0e235..06bfe557 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -9,6 +9,7 @@ import torch +from vllm_rlt.attention.sparse import DENSE, reference_mode from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import Operation, Request from vllm_rlt.worker.buffers import stage_metadata @@ -119,6 +120,9 @@ class PipelineState: position: int tokens: list[PipelineToken] pending: bool = False + # Sparse decode: the request's reference slots and, per pipeline row, the + # index of the slot it reads inside that snapshot. + references: tuple | None = None def to(self, device): """Own hidden copies for suspension; never borrow graph/runner buffers.""" @@ -130,6 +134,9 @@ def to(self, device): ) for token in self.tokens ], + references=None + if self.references is None + else (self.references[0].to(device), self.references[1]), ) @@ -152,6 +159,8 @@ def __init__( *, runtime=None, compute_stream=None, + reference_store=None, + sparse_decode=None, ): self.model, self.cache, self.config = model, cache, config self.metadata = cache.metadata_builder @@ -159,6 +168,9 @@ def __init__( self.pipeline_states: dict[str, PipelineState] = {} self.device = next(model.parameters()).device self.compute_stream = compute_stream + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self.reference_generation = 0 self.submissions = SubmissionQueue() self.input_buffers = { dtype: StagedInputTensor(dtype, self.device, use_uva=execution.use_uva) @@ -175,9 +187,14 @@ def __init__( def _packed(self, packed): capabilities = self.cache.attention_capabilities - return packed and capabilities.packed_prefill and capabilities.packed_matches_decode + return ( + packed + and self.reference_store is None + and capabilities.packed_prefill + and capabilities.packed_matches_decode + ) - def _prepare_inputs(self, batch): + def _prepare_inputs(self, batch, references=()): """Plan all host-known inputs, then publish one snapshot per dtype.""" for buffer in self.input_buffers.values(): buffer.reset() @@ -219,8 +236,32 @@ def metadata(ids, positions, depths, *, packed=False): for active, layouts in steps ] verification = tuple(self.metadata.bind(kv, uploaded) for kv in verification) + if references: + steps = [ + ( + active, + self._reference_metadata(layouts, [(references[i], offset) for i in active]), + ) + for offset, (active, layouts) in enumerate(steps) + ] + owners = [(lease, row) for lease in references for row in range(len(lease))] + verification = self._reference_metadata(verification, owners) return tokens, steps, verification + def _reference_metadata(self, layouts, owners): + result = [] + for kv in layouts: + modes = [reference_mode(depth) for _, depth, _ in kv.rows] + if all(mode == DENSE for mode in modes): + result.append(kv) + continue + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode.key_retention + ) + self.reference_store.bind_positions(plan, kv.position_ids) + result.append(replace(kv, reference_plan=plan)) + return tuple(result) + def _recurrent(self, hidden, metadata): forward = ForwardBatch( Operation.RECURRENT, hidden=hidden, metadata=metadata, graphs=self.graphs @@ -257,11 +298,31 @@ def submit(self, batch): if self.compute_stream is not None else nullcontext() ): - if self.config.mode == "pipeline": - submission = self._submit_pipeline(batch) - else: - inputs, steps, verification = self._prepare_inputs(batch) - submission = self._submit_round(batch, inputs, steps, verification) + references = [] + try: + if self.reference_store is not None: + self.reference_generation += 1 + for index, span in enumerate(batch.requests): + references.append( + self.reference_store.reserve( + span.request.request_id, + index, + span.token_count, + self.reference_generation, + ) + ) + if self.config.mode == "pipeline": + submission = self._submit_pipeline(batch) + else: + inputs, steps, verification = self._prepare_inputs(batch, references) + submission = self._submit_round(batch, inputs, steps, verification) + finally: + if references: + # Every candidate has finished its final recurrent read. + # No reference survives this block round, including the + # rejected suffix. Clear/reuse is ordered on core after + # those readers, without waiting for host acceptance. + self.reference_store.release(references) self.submissions.complete(submission.event) return submission except Exception: