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/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/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_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py new file mode 100644 index 00000000..e7f85e59 --- /dev/null +++ b/tests/test_sparse_decode_attention.py @@ -0,0 +1,739 @@ +"""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, +) +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) + + +# --------------------------------------------------------------------------- +# 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_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_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..663060ca --- /dev/null +++ b/tests/test_sparse_decode_batched.py @@ -0,0 +1,75 @@ +"""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]) +def test_batched_ragged_gqa_matches_scalar_oracle(device, dtype, retention, ties): + 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) + for _ 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 + ) + results.append(implementation(queries, keys, values, tables, lengths, plan, layer=0)) + torch.testing.assert_close(results[1], results[0], rtol=0, atol=0) + for name in ("selected", "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) 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/sparse.py b/vllm_rlt/attention/sparse.py new file mode 100644 index 00000000..4701efc0 --- /dev/null +++ b/vllm_rlt/attention/sparse.py @@ -0,0 +1,436 @@ +"""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 dataclasses import dataclass + +import torch + +# 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) + + +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, +) -> 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. + """ + 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) + 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) + 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, + ) + 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. + 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 +): + """Use bounded row batches on CUDA, retaining the scalar numerical oracle.""" + implementation = batched_paged_attention if q.is_cuda else reference_paged_attention + return implementation( + q, key_cache, value_cache, block_tables, context_lengths, plan, layer=layer + ) + + +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 +): + """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. + """ + 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() + consume_rows = [r for r, mode in enumerate(plan.modes) if mode == CONSUME and lengths[r]] + counts = {} + if consume_rows: + slots = plan.slot_ids[consume_rows] + counts = dict(zip(consume_rows, 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(): + # At most 16M elements per gathered K/V tensor (or one very long row). + chunk = max(1, min(32, 16 * 1024 * 1024 // (extent * heads * dim))) + for start in range(0, len(all_rows), chunk): + rows = all_rows[start : start + chunk] + row_ids = torch.tensor(rows, dtype=torch.int64, device=q.device) + slots = plan.slot_ids[row_ids] + queries = q[row_ids].float() + tables = block_tables[row_ids] + batch_index = torch.arange(len(rows), device=q.device)[:, None, None] + if mode == CONSUME: + columns = torch.arange(extent, device=q.device)[None, None, :] + selected = store.selected[slots[:, None, None], layer, head_index, columns].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 = ( + store.output[slots, layer] + - store.contribution[slots, layer] + + store.mass[slots, layer].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 + keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + 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) + capacity = store.selected.shape[-1] + if count > capacity: + raise ValueError( + f"reference capacity {capacity} cannot hold {count} retained keys" + ) + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + selected = order[..., :count].sort(dim=-1).values + gathered = probabilities.gather(-1, selected) + selected_values = values[batch_index, selected, head_index] + contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) + columns = torch.arange(count, device=q.device)[None, None, :] + # Explicit indices write only the live prefix, never gather the + # entire capacity-sized selected tensor to obtain a prefix view. + 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 + 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..3f926b7c 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -190,3 +190,35 @@ 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 + + def __post_init__(self): + if type(self.enabled) is not bool: + raise ValueError("enabled must be a boolean") + 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..06b6a9d3 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 @@ -719,6 +720,20 @@ def _attend_prepared( self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) if isinstance(batch, DeviceKVMetadata): + plan = getattr(batch, "reference_plan", None) + if plan is not None and not plan.dense(): + # Reference sparse decode path: it reads the same paged tensors + # and publishes/consumes retained state through the plan's + # store. Optimized kernels consume the same plan later. + return sparse_paged_attention( + q, + self.key_cache[:, layer], + self.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + ) return self.attention( q, self.key_cache[:, layer], diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..5dfd5941 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, ): @@ -118,6 +123,80 @@ 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: + # This PR covers the asynchronous self-speculation cell; other + # integration rows are explicit capability failures for now. + if not self.async_speculative: + raise ValueError( + "sparse decode attention currently requires asynchronous self-speculation" + ) + if speculative_config.mode != "block": + raise ValueError("sparse decode attention currently requires block speculation") + # 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" + ) + # Two in-flight rounds may each hold every verification row, so the + # pool covers both rounds' live tokens with room for slot reuse. + slots = ( + 2 * scheduler_config.max_num_seqs * (speculative_config.num_speculative_tokens + 1) + ) + # 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 + * 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, + device=parameter.device, + ) + logger.info( + "sparse decode attention: %d reference slots x %d keys x %d layers x %d heads " + "= %.2f GiB (context bound %d, retention %.2f)", + slots, + capacity, + config.num_hidden_layers, + config.num_attention_heads, + self.reference_store.bytes / 1024**3, + context_bound, + sparse_decode.key_retention, + ) + scratch += self.reference_store.bytes + 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 ) @@ -244,7 +323,13 @@ 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 self.executor = Executor(self.model_runner, speculation=self.speculative_runner) self._pending_speculative = [] self._retiring_speculative = {} 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/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 6b708df2..563fc082 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 DENSE, 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 @@ -50,6 +51,9 @@ 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 class SpeculativeRoundBuffer: @@ -92,6 +96,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 +156,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 +169,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 +186,8 @@ def __init__( copy_stream, rng: RngRegistry | None = None, runtime=None, + reference_store=None, + sparse_decode=None, ): super().__init__( model, @@ -180,6 +198,12 @@ 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) @@ -222,15 +246,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( @@ -264,7 +297,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 +315,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 +401,27 @@ 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 _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 _draft(self, plan, tables, allocations): + 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 +448,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 +469,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 +505,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 +545,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: + 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 +606,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): @@ -685,6 +796,12 @@ 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() 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..9fad5767 --- /dev/null +++ b/vllm_rlt/worker/attention_reference.py @@ -0,0 +1,352 @@ +"""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 dataclasses import dataclass + +import torch + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan + +__all__ = ["ReferenceLease", "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) + + +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): + 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") + 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 + shape = (slots, layers, heads) + self.selected = torch.zeros((*shape, key_bucket or 1), 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(slots)) + self._bindings: dict[int, tuple[str, int, int]] = {} + self._incarnations: dict[tuple[str, 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 outstanding(self) -> int: + """Number of slots currently bound to live tokens.""" + return len(self._bindings) + + 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) + self.incarnation[slot] = count + incarnations.append(count) + return ReferenceLease( + request_id=request_id, + generation=generation, + request_index=request_index, + slots=slots, + incarnations=tuple(incarnations), + slot_ids=torch.tensor(slots, dtype=torch.int64, device=self.device), + 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._clear(torch.tensor(cleared, dtype=torch.int64, device=self.device)) + + def discard(self) -> None: + """Free every binding without device work; the caller has drained.""" + self._bindings.clear() + self.free = list(range(self.slots)) + + 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 lease, row in owners: + 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=torch.tensor(slots, dtype=torch.int64, device=self.device), + store=self, + ) + + def _require_live(self, plan: SparseDecodePlan) -> 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") + for slot, incarnation in zip(plan.slots, plan.incarnations): + binding = self._bindings.get(slot) + if binding is None or binding[2] != incarnation: + raise RuntimeError("stale attention reference plan for a retired token") + + 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 + self.valid[plan.slot_ids[rows], layer] = True + + 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[plan.slot_ids[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") + self.position[plan.slot_ids] = 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] + capacity = self.selected.shape[-1] + if count > capacity: + raise ValueError(f"reference capacity {capacity} cannot hold {count} retained keys") + self.selected[slot, layer, :, :count] = selected.to(torch.int32) + self.counts[slot, layer] = counts.to(torch.int32) + 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 = int(self.counts[slot, layer].max()) + if count == 0: + 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) + rejected = slot_ids[rows > limit] + if rejected.numel(): + # Scalar fills touch only the rejected rows and allocate nothing. + self.valid[rejected] = False + self.counts[rejected] = 0 + self.mass[rejected] = 0 + + 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(): + self.valid[slots] = False + self.uses[slots] = 0 + self.accepted[slots] = 0 + self.incarnation[slots] = 0 + self.position[slots] = 0 + self.mass[slots] = 0 + self.counts[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 + slots = torch.as_tensor(slots, dtype=torch.int64, device=self.device) + self.valid[slots] = False + self.mass[slots] = 0 + self.output[slots] = 0 + self.contribution[slots] = 0 + self.selected[slots] = 0 + self.counts[slots] = 0 diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index 8f2c020b..4c7e7cc0 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -62,6 +62,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 +104,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 +120,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