diff --git a/README.md b/README.md index 00ed4e17..51405e7a 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,10 @@ currently **same-host CUDA** with supported Triton/FA4 paths; it is not a claim of universal multi-node or model support. Speculative steps are not included in mixed prefill/decode batches. + + ### Designed for recurrent inference - **Depth-aware paged KV cache:** LAST_EXITED and SHARED layouts, with optional @@ -92,6 +94,14 @@ in mixed prefill/decode batches. CUDA Graphs, and optional asynchronous/speculative paths. An optional native CUDA SiLU-and-Mul operator accelerates the Ouro activation path, with a Torch fallback; it is not a general model-speedup claim. +- **Sparse decode reference integration (opt-in):** cached loop-2 attention + correction for ordinary decode and block self-speculation, with shared + per-token reference ownership and explicit capability limits. The cross-loop + sparse attention mechanism comes from FlashLoop + ([paper](https://arxiv.org/abs/2609.29812), + [reference implementation](https://github.com/Superone77/FlashLoop)); this + repository contributes the kernel and serving integration. See the + [configuration and scope](docs/runtime.md#sparse-decode-attention-reference-path). - **Offline and online use:** Python API, CLI, and a streaming OpenAI-compatible **completions** endpoint. This is not a chat-completions API. diff --git a/docs/runtime.md b/docs/runtime.md index 460c5e9c..5d61e631 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -230,6 +230,67 @@ The same flags work with the single-device server. A streaming event may contain multiple committed tokens. See [RFC #43](https://github.com/InferMatrixLab/vllm-rlt/issues/43) for the design. +## Sparse decode attention (reference path) + +`SparseDecodeAttentionConfig` enables the RFC #156 cached-reference correction +for generated-token attention. The cross-loop sparse attention mechanism is from +FlashLoop ([paper](https://arxiv.org/abs/2609.29812), +[reference implementation](https://github.com/Superone77/FlashLoop)); the +configuration, kernels and serving integration here are this repository's. +Prompt prefill remains dense. Loop 2 captures +per-token/per-layer/head selection, mass and contributions; loops 3 and 4 reuse +that reference with their current selected KV. Each new input token establishes +its own reference. + +```python +from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig, SpeculativeConfig + +sparse = SparseDecodeAttentionConfig(enabled=True, key_retention=0.1) +# Pass sparse_decode_attention_config=sparse to LLM or StageEngine. +# Ordinary decode uses synchronous scheduling and no speculative_config. +# For synchronous block speculation, also pass: +speculation = SpeculativeConfig(num_speculative_tokens=2, draft_loops=2) +execution = ExecutionConfig(async_scheduling=False, cuda_graphs=False) +``` + +This path supports ordinary synchronous decode and synchronous/asynchronous +block self-speculation with LAST_EXITED KV and refill scheduling. The Torch and +Triton adapters declare sparse-reference support; `auto` selects a compatible +adapter and explicit FlashAttention selection fails rather than silently +replacing the requested backend. Dense and capture rows preserve the selected +backend's output and original packed batch geometry. Selection, retained mass +and contributions, and partial-retention correction use the shared PyTorch +reference implementation. The captured output is the value actually entering +the recurrence, stored in FP32 after the backend's output-dtype conversion. +It has host readbacks and per-row launch costs; enabling it does not +promise faster serving. CUDA Graphs, token pipelines, mixed scheduling, +preemption/recovery and ordinary asynchronous decode are rejected explicitly. + +The retained count is `min(N, max(1, ceil(N * key_retention)))`, including the +current position among the N causal keys. `key_bucket` controls buffer capacity, +not the algorithm: it must hold the retained count over the model's positional +limit, which is the only context bound the reference needs. The index buffer +starts small and grows on the owner stream, so a long model context does not +cost its worst case up front; admission never rejects a request for being long. +The pool holds one slot per live ordinary input token, K+1 per synchronous +block request, and two rounds of K+1 per asynchronous block request. All +reference payload bytes contribute to automatic KV memory planning; explicit +KV block overrides still require the caller to budget total peak memory. + +Ordinary references retire at CODA or request teardown after their last +recurrent reader. Synchronous block references retire on the compute stream +when every verification input has completed its final read; neither accepted +nor rejected rows carry a reference into the next round. With draft depth 1, +verification captures loop 2; with draft depth 2, drafting captures it. +Full retention uses the same dense backend at later loops as a control, while +still validating and consuming the captured reference lifecycle. The standalone +Torch arithmetic oracle is tested separately; different attention algorithms +need not round to identical BF16 values. The +[numerical diagnosis](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029) +is kept as a release attachment rather than a checked-in report. Retention below +1 is approximate and +needs task-quality evaluation before a deployment recommendation. + ## SHARED prefill semantics SHARED retains one physical KV plane and completes every loop for a prompt diff --git a/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_async_speculative.py b/tests/test_async_speculative.py index 4d6d1871..eb45029e 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -173,7 +173,7 @@ def forbidden(*args, **kwargs): @pytest.mark.gpu @pytest.mark.parametrize("graphs", [False, True]) -def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, graphs): +def test_eos_discards_queued_round_and_the_id_frees_after_retirement(monkeypatch, graphs): model = tiny("cuda") e = make_engine(model, graphs=graphs) bootstrap(e, ignore_eos=False) @@ -191,7 +191,17 @@ def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, gra e.step() out = e._collect_speculative()[0] assert out.finished and out.token_ids[-1] == 0 and len(out.token_ids) == 2 - # Reuse while an old ticket still exists. It must never touch the new request. + # The engine never reuses an id while the previous token may still own + # resources. Whether the round has already retired by now depends on the + # graph setting, so assert the rejection only while the engine still lists + # the id; the live-duplicate case is covered by test_stage_engine. + if "r" in e._retiring_speculative: + with pytest.raises(ValueError, match="duplicate request ID"): + e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) + # Once retirement completes the id is free, and the new request never sees + # the old round's state. + e._retire_speculative(wait=True) + assert "r" not in e._retiring_speculative monkeypatch.setattr(e.speculative_runner, "_coda", coda) e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) output = drain(e)["r"] @@ -345,6 +355,8 @@ def test_active_page_tables_are_contiguous_and_do_not_touch_inactive_storage(): assert cache.allocate("short", 4) and cache.allocate("long", 10) runner = object.__new__(AsyncSpeculativeRunner) runner.cache = cache + runner.reference_store = None + runner.sparse_decode_config = None buffer = SimpleNamespace( leased=False, host_tables=torch.full((128,), -91, dtype=torch.int32), diff --git a/tests/test_async_speculative_lifecycle.py b/tests/test_async_speculative_lifecycle.py index b3d469f7..a68ae940 100644 --- a/tests/test_async_speculative_lifecycle.py +++ b/tests/test_async_speculative_lifecycle.py @@ -17,9 +17,9 @@ from vllm_rlt.request import RequestPhase from vllm_rlt.serving.worker import EngineWorker from vllm_rlt.worker.async_speculative import ( + AsyncSpeculativeOutput, AsyncSpeculativeRunner, SpeculativeRequestState, - SpeculativeRoundTicket, ) @@ -105,7 +105,7 @@ def prepared_engine(*, bad_second=False, first_frontier=4, retirement_ready=True host_result=host, ) tickets.append( - SpeculativeRoundTicket( + AsyncSpeculativeOutput( SchedulerOutput(RequestPhase.SPECULATIVE, [scheduled_request]), buffer, (state,), @@ -382,7 +382,7 @@ def test_ticket_result_format_does_not_depend_on_statistics(mode, with_stats): host_result=torch.tensor(values), copy_done=Completion(), leased=True, keepalive=[] ) stats = SpeculativeStats() if with_stats else None - ticket = SpeculativeRoundTicket(batch, buffer, (), (), 2, mode=mode, stats=stats) + ticket = AsyncSpeculativeOutput(batch, buffer, (), (), 2, mode=mode, stats=stats) result = ticket.collect() assert result.speculative[0].token_ids == [11] assert result.speculative[0].accepted_count == 0 diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index e6179d0d..57735bcc 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -53,6 +53,30 @@ def test_cpu_auto_reports_torch_and_runs_reference_attention(caplog): torch.testing.assert_close(selected(q, k, v, table, lengths), explicit(q, k, v, table, lengths)) +def test_sparse_auto_filters_flash_and_explicit_flash_fails(packages, monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0)) + required = BackendCapabilities(sparse_decode=True) + selected = create_backend( + "auto", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + assert selected.info["selected_backend"] == "triton" + assert "sparse_decode" in selected.info["selection_reason"] + with pytest.raises(ValueError, match="lacks required capabilities: sparse_decode"): + create_backend( + "flash_attn_4", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + + @pytest.mark.parametrize( "sm,dtype,head_dim,block_size,missing,expected", [ diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py new file mode 100644 index 00000000..d51b307f --- /dev/null +++ b/tests/test_sparse_decode_attention.py @@ -0,0 +1,862 @@ +"""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 lease identity, delayed +completion and reference-slot reuse. +""" + +import math + +import pytest +import torch + +from tests.helpers import ( + bootstrap_speculation, + make_cpu_cache, + make_speculative_engine, + tiny_ouro_model, +) +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention import sparse as sparse_module +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + REFERENCE_DEPTH, + SparseDecodePlan, + reference_mode, + reference_statistics, + retained_count, + sparse_correction, + sparse_paged_attention, +) +from vllm_rlt.worker.attention_reference import SparseDecodeReferenceStore +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 SparseDecodeReferenceStore(**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), + owners=(None, "lease-1", "lease-1", "lease-2"), + key_retention=0.5, + ) + assert plan.groups() == {DENSE: [0], CAPTURE: [1, 2], CONSUME: [3]} + assert not plan.dense() + + +def test_retained_count_follows_flashloop_ceil_rule(): + # FlashLoop aeaadee select_decode_keys: k = min(N, max(1, ceil(N * r))). + assert retained_count(13, 0.1) == 2 + assert retained_count(10, 0.1) == 1 + assert retained_count(10, 0.5) == 5 + assert retained_count(9, 0.5) == 5 + assert retained_count(1, 0.1) == 1 + assert retained_count(4, 1.0) == 4 + assert retained_count(0, 0.5) == 0 + with pytest.raises(ValueError, match="key_retention"): + retained_count(4, 1.5) + + +def test_reference_statistics_fail_when_capacity_cannot_hold_retained_keys(): + # Capacity never truncates the algorithm: a count above it fails loudly. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.2, 0.3, 0.5])).reshape(3, 1, 1) + values = torch.ones(3, 1, 1) + with pytest.raises(ValueError, match="cannot hold"): + reference_statistics(q, keys, values, key_retention=0.5, key_bucket=1) + # The same query passes with sufficient capacity. + output, _, _, positions = reference_statistics(q, keys, values, key_retention=0.5, key_bucket=2) + assert positions.shape[1] == 2 + assert output.shape == (1, 1) + + +def test_reference_statistics_match_rfc_worked_example(): + # RFC #156 section 3.2: probs [0.1, 0.3, 0.6], values [10, 20, 30]; + # full output 25; retaining the last two keys gives m=0.9, C_ref=24. + # A later loop with local softmax [0.75, 0.25] over values [20, 40] + # corrects to 25 - 24 + 0.9 * (0.75*20 + 0.25*40) = 23.5. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.1, 0.3, 0.6])).reshape(3, 1, 1) + values = torch.tensor([10.0, 20.0, 30.0]).reshape(3, 1, 1) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=0.5) + assert output.item() == pytest.approx(25.0) + assert mass.item() == pytest.approx(0.9) + assert contribution.item() == pytest.approx(24.0) + assert positions.tolist() == [[1, 2]] + + later_keys = torch.log(torch.tensor([3.0, 1.0])).reshape(1, 2, 1) + later_values = torch.tensor([20.0, 40.0]).reshape(1, 2, 1) + corrected = sparse_correction(q, later_keys, later_values, mass, output, contribution) + assert corrected.item() == pytest.approx(23.5) + + +def test_full_retention_matches_dense_attention(): + torch.manual_seed(7) + length, kv_heads, heads, dim = 9, 2, 4, 8 + q = torch.randn(heads, dim) + keys = torch.randn(length, kv_heads, dim) + values = torch.randn(length, kv_heads, dim) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=1.0) + # Full retention keeps every key, so the correction reproduces the dense + # reference output exactly. + dense = reference_statistics(q, keys, values, key_retention=1.0)[0] + groups = heads // kv_heads + corrected = sparse_correction( + q, + keys.permute(1, 0, 2).repeat_interleave(groups, dim=0), + values.permute(1, 0, 2).repeat_interleave(groups, dim=0), + mass, + output, + contribution, + ) + assert positions.shape[1] == length + assert torch.allclose(corrected, dense, atol=1e-5) + assert torch.allclose(output, dense, atol=1e-5) + + +@pytest.mark.parametrize("retention, expected", [(0.5, 23.5), (1.0, 40.0)]) +def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expected): + store = make_store(slots=1, layers=1, heads=1, head_dim=1, key_bucket=3) + lease = store.reserve("r", 0, 1) + q = torch.ones(1, 1, 1) + keys = torch.zeros(1, 4, 1, 1) + values = torch.zeros_like(keys) + keys[0, :3, 0, 0] = torch.log(torch.tensor([0.1, 0.3, 0.6])) + values[0, :3, 0, 0] = torch.tensor([10.0, 20.0, 30.0]) + tables, lengths = torch.tensor([[0]], dtype=torch.int32), torch.tensor([3]) + capture = store.build_plan([(lease, 0)], [CAPTURE], key_retention=retention) + # No dense callback: exercise the actual paged Torch oracle even at r=1. + output = sparse_paged_attention(q, keys, values, tables, lengths, capture, layer=0) + assert output.item() == pytest.approx(25.0) + keys[0, :3, 0, 0] = torch.log(torch.tensor([1.0, 3.0, 1.0])) + values[0, :3, 0, 0] = torch.tensor([100.0, 20.0, 40.0]) + consume = store.build_plan([(lease, 0)], [CONSUME], key_retention=retention) + output = sparse_paged_attention(q, keys, values, tables, lengths, consume, layer=0) + assert output.item() == pytest.approx(expected) + store.release(lease) + + +# --------------------------------------------------------------------------- +# Store lifecycle (RFC #156 section 5) +# --------------------------------------------------------------------------- + + +def test_reserve_uses_one_slot_per_live_token(): + store = make_store() + lease = store.reserve("r", 0, 3) + assert len(lease) == 3 + assert len(set(lease.slots)) == 3 + assert all(store._bindings[slot] is lease for slot in lease.slots) + 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) + fill_payload(store, lease.slots) + store.release(lease) + for slot in lease.slots: + # Readers are gated by counts and mass. 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.counts[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) + 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.counts[rejected].any()) + rejected_plan = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected_plan, layer=0, row=0) + + +def test_accept_drops_host_counts_so_rejected_rows_cannot_size_a_read(): + store = make_store() + lease = store.reserve("r", 0, 3) + fill_payload(store, lease.slots) + for slot in lease.slots: + for layer in range(store.layers): + store.note_count(slot, layer, 1) + store.accept([lease], torch.tensor([1])) + assert all(store.host_count(slot, 0) is None for slot in lease.slots) + # The kept row still reads its device count; the rejected row cannot be read. + kept = store.build_plan([(lease, 1)], [CONSUME], key_retention=0.5) + assert store.read_reference(kept, layer=0, row=0)[0].shape[-1] == 1 + rejected = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected, layer=0, row=0) + + +def test_accept_keeps_root_when_every_candidate_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 3) + 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_assert_liveness_without_state(): + # Availability comes from the retained count and ownership from the lease, so + # capture and consume record nothing; they reject a plan whose slots moved on. + store = make_store() + lease = store.reserve("r", 0, 2) + plan = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + before = store.counts.clone() + store.capture(plan, layer=0) + store.capture(plan, layer=1) + store.consume(plan, layer=0) + assert torch.equal(store.counts, before) + store.release(lease) + with pytest.raises(RuntimeError, match="stale"): + store.capture(plan, layer=0) + + +def test_failed_reservation_keeps_the_pool_intact(monkeypatch): + # A device allocation can fail between taking slots and returning a lease. + # The slots must go back to the pool, or the next request starves. + store = make_store(slots=2) + original = store._index + attempts = {"n": 0} + + def failing(values): + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("injected allocation failure") + return original(values) + + monkeypatch.setattr(store, "_index", failing) + with pytest.raises(RuntimeError, match="injected allocation failure"): + store.reserve("a", 0, 2) + assert len(store.free) == 2 + assert store.outstanding == 0 + + lease = store.reserve("b", 0, 2) + assert len(lease.slots) == 2 + assert len(store.free) == 0 + + +def test_stale_plan_after_a_new_lease_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 2) + stale = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + store.release(lease) + # Same request and rows, a fresh lease. + renewed = store.reserve("r", 0, 2) + assert renewed is not lease + 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) + second = store.reserve("b", 1, 2) + assert not set(first.slots) & set(second.slots) + with pytest.raises(RuntimeError, match="no retired attention reference slots"): + store.reserve("c", 2, 1) + store.release(first) + third = store.reserve("c", 2, 1) + assert set(third.slots) <= set(first.slots) + + +def test_a_reused_slot_is_owned_by_the_new_lease(): + store = make_store() + lease = store.reserve("r", 0, 1) + slot = lease.slots[0] + assert store._bindings[slot] is lease + store.release(lease) + again = store.reserve("r", 0, 1) + assert again.slots[0] == slot + assert again is not lease + assert store._bindings[slot] is again + + +def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): + # A plan must identify the physical slot's current owner. A counter keyed by + # (request_id, row) restarts for the next request, so the previous owner's + # plan could carry the same number, pass the liveness check, and overwrite + # the new owner's reference state. Binding the slot to the lease object makes + # the identity exact. + store = make_store(slots=1, layers=1, heads=1, head_dim=2, key_bucket=4) + first = store.reserve("a", 0, 1) + stale = store.build_plan([(first, 0)], [CAPTURE], key_retention=1.0) + store.release(first) + + second = store.reserve("b", 0, 1) + assert second.slots == first.slots + current = store.build_plan([(second, 0)], [CAPTURE], key_retention=1.0) + assert current.owners[0] is second and stale.owners[0] is first + + store.capture(current, layer=0) + counts = store.counts.clone() + mass = store.mass.clone() + + with pytest.raises(RuntimeError, match="stale"): + store.capture(stale, layer=0) + with pytest.raises(RuntimeError, match="stale"): + store.write_reference( + stale, + 0, + 0, + selected=torch.tensor([[0, 1]]), + counts=torch.tensor([2]), + mass=torch.tensor([0.5]), + output=torch.zeros(1, 2), + contribution=torch.zeros(1, 2), + ) + assert torch.equal(store.counts, counts) + assert torch.equal(store.mass, mass) + + +def test_bytes_matches_layout(): + store = make_store(slots=3, layers=2, heads=2, head_dim=4, key_bucket=8) + # selected (int32 x key_bucket) + counts + mass + output + contribution + expected = 3 * 2 * 2 * 8 * 4 + 3 * 2 * 2 * 4 + 3 * 2 * 2 * 4 + 2 * 3 * 2 * 2 * 4 * 4 + assert store.bytes == expected + + +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), + owners=(object(), object()), + key_retention=0.5, + ) + 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 SparseDecodeReferenceStore(**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) + sparse = decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=1.0, + ) + assert store.outstanding > 0, "the sparse path never held 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) + 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) + 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])) + + store.accept = recording + e.step() + e.step() + e._collect_speculative() + assert observed, "acceptance-driven retirement never ran" + for count, masses in observed: + assert count == accepted + 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_sparse_engine_matches_dense_tokens_at_full_retention(): + # End-to-end through the serving engine: full retention must reproduce the + # dense engine's tokens exactly, including async speculative rounds. + def run(sparse): + model = tiny("cuda") + engine = make_speculative_engine(model, sparse=sparse) + prompts = [[2, 3, 4], [5, 6]] + for i, prompt in enumerate(prompts): + engine.add_request( + str(i), prompt, SamplingParams(max_tokens=12 - 4 * i, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + outputs[output.request_id] = output + tokens = [outputs[str(i)].token_ids for i in range(len(prompts))] + outstanding = engine.reference_store.outstanding if engine.reference_store else 0 + engine.close() + return tokens, outstanding + + dense, _ = run(None) + sparse, outstanding = run(SparseDecodeAttentionConfig(enabled=True, key_retention=1.0)) + assert sparse == dense + assert outstanding == 0 + + +@pytest.mark.gpu +def test_release_from_another_stream_builds_its_index_on_the_owner_stream(monkeypatch): + # A speculative ticket retires from the host's current stream while the + # store clears on its owner stream. An index built on the current stream is + # freed when release returns and can be reused there before the owner + # stream reads it, clearing some other slot. + owner = torch.cuda.Stream() + store = SparseDecodeReferenceStore( + slots=8, layers=1, heads=1, head_dim=4, key_bucket=4, device="cuda", stream=owner + ) + lease = store.reserve("r", 0, 2) + other = next(slot for slot in range(store.slots) if slot not in lease.slots) + torch.cuda.synchronize() + clear, streams = store._clear, [] + + def record(indices): + streams.append(torch.cuda.current_stream()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + with torch.cuda.stream(owner): + torch.cuda._sleep(20_000_000) + store.release(lease) + poison = [torch.full((1,), other, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert streams == [owner] and len(poison) == 32 + assert not any(slot in store._bindings for slot in lease.slots) + assert other not in lease.slots + + +@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 = SparseDecodeReferenceStore( + slots=16, layers=24, heads=16, head_dim=128, key_bucket=32768, device="cuda" + ) + lease = store.reserve("r", 0, 16) + 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_sizes_reference_capacity_from_the_model_context(): + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0) + ) + # Capacity is the worst-case retained count over the model's positional + # limit; the buffer itself starts small and grows on demand. + assert e.reference_store.key_bucket == model.config.max_position_embeddings + assert e.reference_store.selected.shape[-1] < e.reference_store.key_bucket + e.close() + + +@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 + + +def test_store_close_drops_the_device_scratch(monkeypatch): + released = [] + monkeypatch.setattr( + "vllm_rlt.worker.attention_reference.release_workspace", + lambda device: released.append(device), + ) + store = make_store(slots=2) + store.reserve("r", 0, 2) + store.close() + # Teardown must leave nothing bound and hand the kernel scratch back; the + # caller has drained the stream, so the tensors cannot be in use. + assert store.outstanding == 0 + assert not store._bindings + assert released == [store.device] + + +def test_release_workspace_is_scoped_to_the_closed_device(): + from vllm_rlt.kernels.sparse_workspace import ( + CAPTURE_SCRATCH, + CONSUME_SCRATCH, + RETENTION_SCALARS, + release_workspace, + ) + + cpu, other = torch.device("cpu"), torch.device("meta") + CAPTURE_SCRATCH[(cpu, 1)] = torch.zeros(4) + CONSUME_SCRATCH[(cpu, 1)] = torch.zeros(4) + RETENTION_SCALARS[(cpu, 0.5)] = torch.zeros(1, dtype=torch.float64) + CAPTURE_SCRATCH[(other, 2)] = torch.zeros(4) + try: + release_workspace(cpu) + assert (cpu, 1) not in CAPTURE_SCRATCH + assert (cpu, 1) not in CONSUME_SCRATCH + assert (cpu, 0.5) not in RETENTION_SCALARS + # A second engine on another device keeps its scratch. + assert (other, 2) in CAPTURE_SCRATCH + finally: + release_workspace() + assert not CAPTURE_SCRATCH and not CONSUME_SCRATCH and not RETENTION_SCALARS diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py new file mode 100644 index 00000000..bb204f60 --- /dev/null +++ b/tests/test_sparse_decode_batched.py @@ -0,0 +1,265 @@ +"""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, + slot_span, +) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("use_dense_callback", [False, True]) +def test_batched_ragged_gqa_matches_scalar_oracle( + device, dtype, retention, ties, use_dense_callback +): + torch.manual_seed(157) + rows, heads, kv_heads, dim, block_size = 6, 4, 2, 128, 16 + lengths = torch.tensor([0, 17, 17, 31, 31, 257], device=device) + tables = torch.randperm(rows * 17, device=device).reshape(rows, 17).int() + keys = torch.randn(rows * 17, block_size, kv_heads, dim, device=device, dtype=dtype) + values = torch.randn_like(keys) + queries = torch.randn(rows, heads, dim, device=device, dtype=dtype) + if ties: + queries.zero_() + stores = [ + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=257, + device=device, + initial_key_capacity=1 if implementation else None, + ) + for implementation in range(2) + ] + leases = [store.reserve("test", 0, rows) for store in stores] + implementations = [reference_paged_attention, batched_paged_attention] + for modes in ([CAPTURE] * rows, [DENSE, CAPTURE, CONSUME, CONSUME, DENSE, CONSUME]): + results = [] + for store, lease, implementation in zip(stores, leases, implementations): + plan = store.build_plan( + [(lease, row) for row in range(rows)], modes, key_retention=retention + ) + # A deliberately distinct backend output verifies that capture + # stores the rounded recurrence input, not its FP32 oracle PV. + callback = (lambda q, *args: q * 0.375 + 1.25) if use_dense_callback else None + results.append( + implementation( + queries, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=callback, + ) + ) + torch.testing.assert_close(results[1], results[0], rtol=0, atol=0) + capacity = stores[1].selected.shape[-1] + assert torch.equal(stores[0].selected[..., :capacity], stores[1].selected) + assert not torch.count_nonzero(stores[0].selected[..., capacity:]) + for name in ("counts",): + 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) + 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) + with pytest.raises(RuntimeError, match="stale attention reference"): + batched_paged_attention(*args, layer=0) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_reference_growth_preserves_live_layers_and_budget(device): + store = make_store( + slots=3, + layers=2, + heads=2, + head_dim=4, + key_bucket=17, + initial_key_capacity=2, + device=device, + ) + leases = [store.reserve(str(i), i, 1) for i in range(2)] + plans = [store.build_plan([(lease, 0)], [CAPTURE], key_retention=1.0) for lease in leases] + # Simulate queued work on the runner's owner stream before reallocating. + if device == "cuda": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + store.stream = stream + context = torch.cuda.stream(stream) + else: + from contextlib import nullcontext + + context = nullcontext() + budget = store.peak_capacity_bytes + fixed = store.bytes - store.selected.nbytes + assert budget == fixed + 3 * 2 * 2 * (2 + 4 + 8 + 16 + 17) * 4 + with context: + for layer in range(2): + store.write_reference( + plans[0], + layer, + 0, + selected=torch.tensor([[1, 3], [2, 4]], device=device), + counts=torch.full((2,), 2, device=device), + mass=torch.full((2,), 0.5, device=device), + output=torch.full((2, 4), 0.25, device=device), + contribution=torch.full((2, 4), 0.125, device=device), + ) + before = [x.clone() for x in store.read_reference(plans[0], 1, 0)] + # Grow several times while a different lease remains live. + for count in [3, 7, 17]: + store.ensure_key_capacity(count) + assert store.selected.shape[-1] >= count + assert store.peak_capacity_bytes == budget + for actual, expected in zip(store.read_reference(plans[0], 1, 0), before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + with pytest.raises(ValueError, match="capacity"): + store.ensure_key_capacity(18) + if device == "cuda": + store.stream.synchronize() + for lease in leases: + store.release(lease) + assert store.outstanding == 0 + + +@pytest.mark.gpu +def test_reference_growth_rejects_foreign_stream(): + store = make_store(key_bucket=17, initial_key_capacity=2, device="cuda") + store.stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="owner stream"): + store.ensure_key_capacity(3) + assert store.selected.shape[-1] == 2 + + +def test_slot_span_requires_every_slot_in_place(): + # Endpoint arithmetic accepts [5, 3, 7] (5 + 3 == 7 + 1) and would then use + # the slice [5, 6, 7], which is another request's reference. + assert slot_span([5, 3, 7]) is None + assert slot_span([7, 6, 5]) is None + assert slot_span([5, 6, 8]) is None + assert slot_span([5, 5, 5]) is None + assert slot_span([3, 4, 5]) == slice(3, 6) + assert slot_span([5]) == slice(5, 6) + + +def _fragmented_store(): + """A store whose next three-row reservation yields slots ``[5, 3, 7]``. + + ``reserve`` reverses the tail of the free list, so releasing 7, then 3, then + 5 leaves ``[7, 3, 5]`` and the reservation comes back fragmented. Slots 6 and + 4 stay bound to unrelated live leases, which is what a span read of + ``[5, 6, 7]`` would corrupt. + """ + store = make_store(slots=8, layers=1, heads=2, head_dim=4, key_bucket=8) + held = [store.reserve(f"held{i}", i, 1) for i in range(5)] # slots 7, 6, 5, 4, 3 + store.release([held[0], held[4], held[2]]) + assert store.free[-3:] == [7, 3, 5] + lease = store.reserve("fragmented", 0, 3) + assert lease.slots == (5, 3, 7) + return store, lease, held[1], held[3] + + +def test_fragmented_consume_reads_each_rows_own_reference(): + torch.manual_seed(157) + store, lease, live_six, live_four = _fragmented_store() + rows, heads, kv_heads, dim, block_size, extent = 3, 2, 2, 4, 8, 8 + # A zero query makes attention uniform, so the correction is identical for + # every row and each output equals its own slot's reference plus one + # constant. Reading the span [5, 6, 7] breaks exactly one row, by 3. + queries = torch.zeros(rows, heads, dim) + keys = torch.randn(1, block_size, kv_heads, dim) + values = torch.randn(1, block_size, kv_heads, dim) + tables = torch.zeros(rows, 1, dtype=torch.int32) + lengths = torch.full((rows,), extent) + + store.output.copy_(torch.arange(8).float()[:, None, None, None]) + store.contribution.zero_() + store.mass.fill_(1.0) + store.counts.fill_(extent) + store.selected.copy_(torch.arange(extent).expand(8, 1, heads, extent)) + for slot in lease.slots: + store.note_count(slot, 0, extent) + + plan = store.build_plan( + [(lease, 0), (lease, 1), (lease, 2)], [CONSUME] * rows, key_retention=0.5 + ) + output = batched_paged_attention(queries, keys, values, tables, lengths, plan, layer=0).float() + + own_reference = torch.tensor([5.0, 3.0, 7.0])[:, None, None] + correction = output - own_reference + torch.testing.assert_close(correction, correction[0].expand_as(correction), rtol=0, atol=1e-6) + # The unrelated owners keep their payload. + assert float(store.output[6, 0].mean()) == 6.0 + assert float(store.output[4, 0].mean()) == 4.0 + store.release([lease, live_six, live_four]) + + +def test_fragmented_capture_writes_only_the_slots_it_reserved(): + torch.manual_seed(157) + store, lease, live_six, live_four = _fragmented_store() + rows, heads, kv_heads, dim, block_size, extent = 3, 2, 2, 4, 8, 8 + queries = torch.randn(rows, heads, dim) + keys = torch.randn(1, block_size, kv_heads, dim) + values = torch.randn(1, block_size, kv_heads, dim) + tables = torch.zeros(rows, 1, dtype=torch.int32) + lengths = torch.full((rows,), extent) + + # A sentinel per slot, so a write to the wrong slot is visible. + sentinel = torch.arange(8).float() - 1000.0 + store.output.copy_(sentinel[:, None, None, None]) + store.contribution.copy_(sentinel[:, None, None, None]) + store.mass.copy_(sentinel[:, None, None]) + store.counts.fill_(-1) + + plan = store.build_plan( + [(lease, 0), (lease, 1), (lease, 2)], [CAPTURE] * rows, key_retention=1.0 + ) + batched_paged_attention(queries, keys, values, tables, lengths, plan, layer=0) + + for slot in lease.slots: + assert float(store.output[slot, 0].mean()) != float(sentinel[slot]), slot + assert bool((store.counts[slot, 0] == extent).all()), slot + # Slot 6 belongs to another request: the span [5, 6, 7] overwrote it and + # left slot 3 unwritten. + for slot in (4, 6): + assert float(store.output[slot, 0].mean()) == float(sentinel[slot]), slot + assert float(store.mass[slot, 0].mean()) == float(sentinel[slot]), slot + assert bool((store.counts[slot, 0] == -1).all()), slot + store.release([lease, live_six, live_four]) diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py new file mode 100644 index 00000000..d5adb3cf --- /dev/null +++ b/tests/test_sparse_decode_fused.py @@ -0,0 +1,524 @@ +"""Opt-in fused sparse decode kernels versus the batched oracle-shaped path.""" + +import pytest +import torch + +from tests.helpers import make_speculative_engine, tiny_ouro_model +from tests.test_sparse_decode_attention import drain, make_store +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + reference_paged_attention, + sparse_paged_attention, +) + + +def test_sparse_config_validates_implementation(): + assert SparseDecodeAttentionConfig().implementation == "batched" + assert SparseDecodeAttentionConfig(implementation="fused").implementation == "fused" + with pytest.raises(ValueError, match="implementation"): + SparseDecodeAttentionConfig(implementation="native") + + +def test_fused_store_keeps_the_oracle_on_cpu(monkeypatch): + calls = [] + import vllm_rlt.attention.sparse as sparse + + monkeypatch.setattr(sparse, "reference_paged_attention", lambda *a, **k: calls.append(1)) + store = make_store(slots=1, layers=1, heads=2, head_dim=4, implementation="fused") + lease = store.reserve("cpu", 0, 1) + plan = store.build_plan([(lease, 0)], [CAPTURE], key_retention=0.5) + q = torch.ones(1, 2, 4) + cache = torch.ones(1, 16, 2, 4) + sparse.sparse_paged_attention( + q, cache, cache, torch.zeros(1, 1, dtype=torch.int32), torch.tensor([3]), plan, layer=0 + ) + assert calls == [1] + + +@pytest.fixture(autouse=True) +def _check_fast_launches(monkeypatch): + """Every reused launch must be the kernel Triton's own dispatch would pick.""" + try: + from vllm_rlt.kernels import sparse_decode + except ImportError: # CPU-only installs have no Triton and no fused launches + return + monkeypatch.setattr(sparse_decode, "_LAUNCH_CHECK", True) + + +@pytest.mark.gpu +@pytest.mark.parametrize("context", [37, 1024, 4111]) +def test_fast_launches_match_triton_dispatch_bit_for_bit(context, monkeypatch): + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + if not sparse_decode._FAST_LAUNCH: + pytest.skip("fast launches are limited to Triton 3.8") + rows, heads, kv_heads, dim, block_size, layers = 3, 4, 2, 64, 16, 18 + pages = -(-context // block_size) + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.bfloat16, 3 + ) + lengths = torch.tensor([context, max(1, context - 17), max(1, context // 2)], device="cuda") + results = {} + for fast in (False, True): + monkeypatch.setattr(sparse_decode, "_FAST_LAUNCH", fast) + store = make_store( + slots=rows, + layers=layers, + heads=heads, + head_dim=dim, + key_bucket=context, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows) + owners = [(lease, row) for row in range(rows)] + outputs = [] + # Layers 0, 1, 16 and 17 differ in Triton's integer specialization. + for _ in range(2): + for layer in (0, 1, 2, 16, 17): + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.1) + outputs.append( + sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=layer, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + ) + torch.cuda.synchronize() + results[fast] = ( + outputs, + store.selected.clone(), + store.counts.clone(), + store.mass.clone(), + store.contribution.clone(), + store.output.clone(), + ) + slow, fast = results[False], results[True] + for a, b in zip(slow[0], fast[0]): + assert torch.equal(a, b) + for a, b in zip(slow[1:], fast[1:]): + assert torch.equal(a, b) + + +@pytest.mark.gpu +def test_consume_scratch_follows_the_head_dimension(monkeypatch): + """A stream's consume scratch must not be reused across head dimensions.""" + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + def consume(dim): + rows, heads, block_size, pages = 3, 4, 16, 5 + q, keys, values, tables = _paged(rows, heads, 2, dim, pages, block_size, torch.float32, 4) + lengths = torch.tensor([7, 33, 80], device="cuda", dtype=torch.int32) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=80, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows) + owners = [(lease, row) for row in range(rows)] + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.25) + output = sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + torch.cuda.synchronize() + return output + + monkeypatch.setattr(sparse_decode, "CONSUME_SCRATCH", {}) + fresh = consume(128) + sparse_decode.CONSUME_SCRATCH.clear() + consume(16) + # The same stream now holds 16-wide scratch; a 128-wide consume must regrow it. + assert torch.equal(consume(128), fresh) + + +def _paged(rows, heads, kv_heads, dim, pages, block_size, dtype, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + tables = torch.randperm(rows * pages, device="cuda", generator=generator) + keys = torch.randn( + rows * pages, block_size, kv_heads, dim, device="cuda", dtype=dtype, generator=generator + ) + values = torch.randn(keys.shape, device="cuda", dtype=dtype, generator=generator) + queries = torch.randn(rows, heads, dim, device="cuda", dtype=dtype, generator=generator) + return queries, keys, values, tables.reshape(rows, pages).int() + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.5]) +@pytest.mark.parametrize("ties", [False, True]) +# The Triton backend takes the split capture; any other callable takes the fallback. +@pytest.mark.parametrize("split", [True, False]) +def test_fused_capture_and_consume_track_batched_path(retention, ties, split): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention + + backend = TritonPagedAttention() if split else (lambda *args: triton_paged_attention(*args)) + + rows, heads, kv_heads, dim, block_size, pages = 6, 4, 2, 128, 16, 17 + # 10 and 30 sit on float rounding edges of ceil(length * 0.1) (30 * 0.1 > 3 in float64). + lengths = torch.tensor([0, 10, 17, 30, 200, 257], device="cuda", dtype=torch.int32) + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=272, + device="cuda", + initial_key_capacity=1, + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [store.reserve("fused", 0, rows) for store in stores] + owners = [[(lease, row) for row in range(rows)] for lease in leases] + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.float32, 1 + ) + if ties: + # Every score equal: both paths must keep the smallest positions. + q.zero_() + captured = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CAPTURE] * rows, key_retention=retention) + captured.append( + sparse_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + batched, fused = stores + # Capture rows return (and publish) the dense backend output unchanged. These + # inputs are FP32, where some architectures (sm_90) round the split copy one + # ulp from the dense kernel; the BF16 test below stays bit for bit. + exact = 0 if not split else 1e-6 + torch.testing.assert_close(captured[1], captured[0], rtol=0, atol=exact) + torch.testing.assert_close(fused.output[:, 1], batched.output[:, 1], rtol=0, atol=exact) + torch.testing.assert_close(fused.counts[:, 1], batched.counts[:, 1], rtol=0, atol=0) + # Both paths leave the same slots bound to their leases; availability is + # decided by the retained count, which the comparison above already covers. + assert set(fused._bindings) == set(batched._bindings) + for row, length in enumerate(lengths.tolist()): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + count = int(batched.counts[slot_b, 1, 0]) + for head in range(heads): + # Both paths write ascending positions; ties keep the smaller ones. + expected = batched.selected[slot_b, 1, head, :count].tolist() + assert fused.selected[slot_f, 1, head, :count].tolist() == expected + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + # A later loop reads its own plane: fresh keys and queries, same pages. + q2, keys2, values2, _ = _paged(rows, heads, kv_heads, dim, pages, block_size, torch.float32, 2) + consumed = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CONSUME] * rows, key_retention=retention) + consumed.append( + sparse_paged_attention( + q2, + keys2, + values2, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + torch.testing.assert_close(consumed[1], consumed[0], rtol=1e-4, atol=1e-5) + assert not consumed[1][0].any(), "an empty row stays zero" + assert fused.outstanding > 0 + for store, lease in zip(stores, leases): + store.release(lease) + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_radix_selection_keeps_oracle_ties_on_repeated_scores(retention): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # Keys cycle through three vectors, so every head sees three distinct scores + # repeated across the context: the retained set is decided almost entirely + # by the tie rule (smaller positions first), with negative scores included. + # Small integers keep every dot product exact. With random reals the dense + # Triton kernel can score one key a ULP apart at different tile positions + # (seen on sm_89), so the "ties" would not be ties for either path. + rows, heads, dim, block_size, pages = 3, 4, 64, 16, 40 + generator = torch.Generator(device="cuda").manual_seed(11) + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + q = integers(rows, heads, dim) + lengths = torch.tensor([7, 300, 640], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=640, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + for store in stores: + lease = store.reserve("ties", 0, rows) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + batched, fused = stores + assert torch.equal(fused.counts, batched.counts) + for slot in range(rows): + for head in range(heads): + count = int(batched.counts[slot, 0, head]) + assert ( + fused.selected[slot, 0, head, :count].tolist() + == batched.selected[slot, 0, head, :count].tolist() + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("folded", [True, False]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_parallel_select_matches_batched_across_many_parts( + ties, retention, folded, monkeypatch +): + import vllm_rlt.kernels.sparse_decode as fused_module + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # The first radix pass runs inside the split kernel or as its own pass. + monkeypatch.setattr(fused_module, "_PASS0_IN_SPLIT", folded) + + # Long rows split into up to 16 parts; with repeated integer scores the + # retained ties span part boundaries, so offsets and tie quotas must carry. + rows, heads, dim, block_size, pages = 4, 4, 64, 16, 320 + generator = torch.Generator(device="cuda").manual_seed(29) + if ties: + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + q = integers(rows, heads, dim) + else: + # Mostly distinct but still exact scores: a one-ULP difference between the + # Triton and batched scores at this length could cross the boundary. + keys = torch.randint( + -8, 9, (rows * pages, block_size, heads, dim), device="cuda", generator=generator + ).float() + q = torch.randint(-8, 9, (rows, heads, dim), device="cuda", generator=generator).float() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + lengths = torch.tensor([1, 129, 2049, 5000], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=5120, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [] + for store in stores: + lease = store.reserve("parts", 0, rows) + leases.append(lease) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=1, dense_attention=backend + ) + batched, fused = stores + for row in range(rows): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + assert torch.equal(fused.counts[slot_f, 1], batched.counts[slot_b, 1]) + for head in range(heads): + count = int(batched.counts[slot_b, 1, head]) + assert ( + fused.selected[slot_f, 1, head, :count].tolist() + == batched.selected[slot_b, 1, head, :count].tolist() + ) + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_split_capture_is_the_dense_backend_output_bit_for_bit(dtype): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + rows, heads, kv_heads, dim = 5, 4, 2, 128 + lengths = torch.tensor([0, 1, 33, 515, 1200], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 76, 16, dtype, 5) + backend = TritonPagedAttention() + expected = backend(q, keys, values, tables, lengths) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=1216, + device="cuda", + implementation="fused", + ) + lease = store.reserve("split", 0, rows) + plan = store.build_plan([(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=0.1) + out = sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + # The split copy and the folded merge must not drift from the dense kernels: + # bit for bit in BF16; in FP32 some architectures (sm_90) contract the split + # copy's arithmetic differently and round it one ulp apart. + atol = 0 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(out, expected, rtol=0, atol=atol) + # O_ref is the output actually fed into the recurrence: rounded, then FP32. + live = list(lease.slots)[1:] + torch.testing.assert_close(store.output[live, 0], expected[1:].float(), rtol=0, atol=atol) + assert all(slot in store._bindings for slot in lease.slots) + + +@pytest.mark.gpu +def test_fused_bf16_consume_matches_scalar_oracle_within_rounding(): + from vllm_rlt.kernels.paged_attention import triton_paged_attention + + rows, heads, kv_heads, dim = 4, 4, 4, 64 + lengths = torch.tensor([5, 64, 129, 300], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 19, 16, torch.bfloat16, 3) + results = [] + for implementation, function in ( + ("fused", sparse_paged_attention), + ("batched", reference_paged_attention), + ): + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=304, + device="cuda", + implementation=implementation, + ) + lease = store.reserve("bf16", 0, rows) + owner = [(lease, row) for row in range(rows)] + for modes in ([CAPTURE] * rows, [CONSUME] * rows): + plan = store.build_plan(owner, modes, key_retention=0.1) + out = function( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=triton_paged_attention, + ) + results.append(out.float()) + # Different FP32 reduction order may move a BF16 rounding step. + torch.testing.assert_close(results[0], results[1], rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +def test_fused_consume_without_capture_is_nan_not_silent(): + rows, heads, dim = 2, 2, 64 + store = make_store( + slots=rows, layers=1, heads=heads, head_dim=dim, device="cuda", implementation="fused" + ) + lease = store.reserve("missing", 0, rows) + plan = store.build_plan([(lease, 0), (lease, 1)], [CONSUME, CONSUME], key_retention=0.5) + q, keys, values, tables = _paged(rows, heads, heads, dim, 2, 16, torch.float32, 4) + lengths = torch.tensor([0, 9], device="cuda", dtype=torch.int32) + out = sparse_paged_attention(q, keys, values, tables, lengths, plan, layer=0) + assert not out[0].any() and torch.isnan(out[1]).all() + + +@pytest.mark.gpu +def test_fused_engine_rounds_abort_and_drain(monkeypatch): + import vllm_rlt.kernels.sparse_decode as fused_module + + calls = {CAPTURE: 0, CONSUME: 0} + original = fused_module.fused_paged_attention + + def counting(*args, **kwargs): + calls[args[5].modes[0]] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(fused_module, "fused_paged_attention", counting) + # The engine's Triton backend adapter must reach the split capture path. + monkeypatch.setattr(fused_module, "_capture_kernel", None) + model = tiny_ouro_model("cuda") + engine = make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused"), + ) + store = engine.reference_store + assert store.implementation == "fused" + params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in (5, 9, 50)] + for i in range(3): + engine.add_request(str(i), [2, 3, 4], params[i]) + outputs = {} + for _ in range(1000): + if "0" in outputs and outputs["0"].finished: + break + for output in engine.step(): + outputs[output.request_id] = output + if "2" not in outputs or not outputs["2"].finished: + engine.abort_request("2") + outputs.update(drain(engine)) + assert outputs["0"].finished and len(outputs["0"].token_ids) == 5 + assert outputs["1"].finished and len(outputs["1"].token_ids) == 9 + assert calls[CAPTURE] > 0 and calls[CONSUME] > 0 + assert store.outstanding == 0 and len(store.free) == store.slots + assert engine.cache_manager.num_used_blocks == 0 + engine.close() diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py new file mode 100644 index 00000000..9be12932 --- /dev/null +++ b/tests/test_sparse_decode_sync.py @@ -0,0 +1,423 @@ +"""Ordinary and synchronous block decode through the shared sparse contract.""" + +from dataclasses import replace + +import pytest +import torch + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.attention import BackendCapabilities, create_backend +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE +from vllm_rlt.engine.stage_engine import StageEngine + + +def engine(*, draft=None, retention=1.0, device="cpu", dtype=torch.float32, **options): + values = dict( + attention_backend="triton" if device == "cuda" else "torch", + cache_config=CacheConfig(512, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=3 + ), + execution_config=ExecutionConfig(static_buffers=True, pad_to_power_of_two=True), + speculative_config=None if draft is None else SpeculativeConfig(2, draft_loops=draft), + sparse_decode_attention_config=( + None + if retention is None + else SparseDecodeAttentionConfig(enabled=True, key_retention=retention) + ), + ) + values.update(options) + return StageEngine(tiny_ouro_model(device, dtype), **values) + + +def generate(e, count=3): + for i in range(count): + e.add_request( + str(i), [2, 3, 4, 5][: 2 + i % 3], SamplingParams(max_tokens=8, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not e.has_unfinished_requests(): + break + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + assert len(outputs) == count + assert all(len(tokens) == 8 for tokens in outputs.values()) + assert e.cache_manager.num_used_blocks == 0 + if e.reference_store is not None: + assert e.reference_store.outstanding == 0 + return outputs + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_full_retention_matches_dense_and_reclaims_slots(draft): + dense = engine(draft=draft, retention=None) + sparse = engine(draft=draft) + try: + assert generate(sparse) == generate(dense) + finally: + sparse.close() + dense.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_plans_follow_tokens_and_keep_prompt_dense(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + seen = [] + original = e.cache_manager._attend_prepared + + def attend(layer, kv, q): + plan = getattr(kv, "reference_plan", None) + if plan is not None: + assert len(plan.modes) == len(q) + for row, mode in enumerate(plan.modes): + if mode == DENSE: + continue + allocation, depth, position = kv.rows[row] + request_id = next(rid for rid, owner in kv.allocations if owner is allocation) + assert position >= len(e.scheduler.requests[request_id].prompt_token_ids) + assert mode == (CAPTURE if depth == 1 else CONSUME) + assert plan.owners[row].request_id == request_id + seen.append((mode, position, plan.slots[row])) + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", attend) + try: + generate(e) + assert {mode for mode, _, _ in seen} == {CAPTURE, CONSUME} + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cancel_and_request_id_reuse_cannot_reuse_old_reference(draft): + e = engine(draft=draft, retention=0.5) + try: + for _ in range(2): + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=32, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding or ( + draft is not None and e.speculative_runner.stats.rounds + ): + break + e.abort_request("r") + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + finally: + e.close() + + +@pytest.mark.parametrize("setting", ["graphs", "pipeline", "preemption", "shared"]) +def test_unqualified_execution_modes_fail_explicitly(setting): + options = { + "graphs": dict(execution_config=ExecutionConfig(cuda_graphs=True)), + "pipeline": dict(speculative_config=SpeculativeConfig(2, mode="pipeline")), + "preemption": dict(scheduler_config=SchedulerConfig(enable_preemption=True)), + "shared": dict(cache_config=CacheConfig(512, 2, layout="shared")), + }[setting] + with pytest.raises(ValueError, match="sparse decode attention"): + engine(**options) + + +def test_backend_explicitly_declares_sparse_reference_support(): + selected = create_backend( + "auto", + torch.device("cpu"), + torch.float32, + 8, + 2, + required_capabilities=BackendCapabilities(sparse_decode=True), + ) + assert selected.capabilities.sparse_decode + + +@pytest.mark.gpu +@pytest.mark.parametrize("kv_heads", [2, 16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): + from vllm_rlt.core.kv_cache_manager import KVCacheManager + from vllm_rlt.worker.attention_reference import SparseDecodeReferenceStore + + # BF16 rounding of Torch softmax and Triton split-KV can differ even for + # identical operands. Dense/capture rows and the full-retention control + # must preserve the configured backend, including in a mixed plan. + torch.manual_seed(157) + rows, heads, dim, length = 32, 16, 128, 130 + cache = KVCacheManager( + 1, + kv_heads, + dim, + rows * 9, + 16, + 1, + device="cuda", + dtype=torch.bfloat16, + backend="triton", + ) + store = SparseDecodeReferenceStore( + slots=rows, layers=1, heads=heads, head_dim=dim, key_bucket=length, device="cuda" + ) + ids = [str(i) for i in range(rows)] + owners = [] + for rid in ids: + assert cache.allocate(rid, length) + keys = torch.randn(length, kv_heads, dim, device="cuda", dtype=torch.bfloat16) + values = torch.randn_like(keys) + cache.write(0, [rid] * length, [0] * length, list(range(length)), keys, values) + lease = store.reserve(rid, 0, 1) + owners.append((lease, 0)) + batch = cache._prepare_batch(ids, [0] * rows, [length - 1] * rows) + q = torch.randn(rows, heads, dim, device="cuda", dtype=torch.bfloat16) + expected = cache._attend_prepared(0, batch, q) + modes = [DENSE] + [CAPTURE] * (rows - 1) + plan = store.build_plan(owners, modes, key_retention=retention) + actual = cache._attend_prepared(0, replace(batch, reference_plan=plan), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + for row in range(1, rows): + torch.testing.assert_close( + store.output[plan.slots[row], 0], expected[row].float(), atol=0, rtol=0 + ) + if retention == 1.0: + # Changed later-loop Q/K/V prevents an identity-only correction test. + q = torch.randn_like(q) + cache._write_prepared( + 0, batch, torch.randn_like(keys[:rows]), torch.randn_like(keys[:rows]) + ) + expected = cache._attend_prepared(0, batch, q) + consume = store.build_plan( + owners, [DENSE] + [CONSUME] * (rows - 1), key_retention=retention + ) + actual = cache._attend_prepared(0, replace(batch, reference_plan=consume), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + # Consume rows read the reference: their slots stay bound until release. + assert store.outstanding == len(owners) + for lease, _ in owners: + store.release(lease) + for rid in ids: + cache.free(rid) + torch.cuda.synchronize() + assert store.outstanding == cache.num_used_blocks == 0 + + +@pytest.mark.parametrize("draft", [1, 2]) +def test_sparse_verification_matches_serial_same_input_hidden_logits_and_kv(draft): + from tests.helpers import make_cpu_cache + from tests.test_sparse_decode_attention import make_reference_store + from vllm_rlt.attention.sparse import reference_mode + from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput + from vllm_rlt.request import Request, RequestPhase + + e = engine(draft=draft, retention=0.5) + model, cache = e.model, e.cache_manager + prefix, current = [3, 4, 5], 7 + assert cache.allocate("r", 16) + for position, token in enumerate(prefix): + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + hidden, _ = model.recurrent(hidden, ["r"], [depth], [position], cache) + request = Request("r", prefix, SamplingParams(ignore_eos=True), generated_token_ids=[current]) + seen, heads = {}, [] + recurrent = model.recurrent_prepared + + def record(hidden, kv, cache, **kwargs): + out, gate = recurrent(hidden, kv, cache, **kwargs) + for row, (_, depth, position) in enumerate(kv.rows): + seen[depth, position] = out[row].clone() + return out, gate + + model.recurrent_prepared = record + hook = model.lm_head.register_forward_hook( + lambda module, inputs, output: heads.append(output.detach().clone()) + ) + try: + e.speculative_runner.submit( + SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(request, 3, 3)]) + ).collect() + finally: + hook.remove() + model.recurrent_prepared = recurrent + assert e.reference_store.outstanding == 0 + tokens = prefix + [current] + tokens.extend(int(model.coda(seen[draft - 1, pos]).argmax()) for pos in (3, 4)) + oracle = make_cpu_cache(model.config) + assert oracle.allocate("r", 16) + store = make_reference_store(model) + for position, token in enumerate(tokens): + lease = store.reserve("r", 0, 1) + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + kv = oracle._prepare_batch(["r"], [depth], [position]) + mode = reference_mode(depth) + if position >= len(prefix) and mode != DENSE: + plan = store.build_plan([(lease, 0)], [mode], key_retention=0.5) + kv = replace(kv, reference_plan=plan) + hidden, _ = model.recurrent_prepared(hidden, kv, oracle, compute_gate=False) + if position >= len(prefix): + torch.testing.assert_close(seen[depth, position], hidden[0], atol=4e-5, rtol=4e-5) + if position >= len(prefix): + torch.testing.assert_close( + heads[-1][position - len(prefix)], model.coda(hidden)[0], atol=4e-5, rtol=4e-5 + ) + store.release(lease) + for depth in range(4): + for layer in range(model.config.num_hidden_layers): + for actual, expected in zip( + cache.read(layer, "r", depth, len(tokens)), + oracle.read(layer, "r", depth, len(tokens)), + ): + torch.testing.assert_close(actual, expected, atol=4e-5, rtol=4e-5) + cache.free("r") + oracle.free("r") + e.close() + + +@pytest.mark.parametrize("draft", [1, 2]) +@pytest.mark.parametrize("accepted", [0, 1, 2]) +def test_sync_rejection_retires_all_references_before_replacement(draft, accepted, monkeypatch): + e = engine(draft=draft, retention=0.5) + try: + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=16, ignore_eos=True)) + while not e.step(): + pass + request = e.scheduler.requests["r"] + start = request.position + original = e.speculative_runner._coda + calls = 0 + + def controlled(hidden): + nonlocal calls + logits = torch.full((len(hidden), 64), -100.0) + if calls < 2: + logits[:, 10 + calls] = 100 + else: + for row in range(3): + logits[row, 10 + row if row < accepted else 20 + row] = 100 + calls += 1 + return logits + + seen = [] + build = e.reference_store.build_plan + + def record(*args, **kwargs): + plan = build(*args, **kwargs) + seen.append(plan) + return plan + + monkeypatch.setattr(e.reference_store, "build_plan", record) + monkeypatch.setattr(e.speculative_runner, "_coda", controlled) + output = e.step()[0] + expected = [10 + row if row < accepted else 20 + row for row in range(accepted + 1)] + assert output.token_ids[-(accepted + 1) :] == expected + assert request.position == start + accepted + 1 + assert e.reference_store.outstanding == 0 + old = seen[0] + with pytest.raises(RuntimeError, match="stale"): + e.reference_store.read_reference(old, 0, 0) + monkeypatch.setattr(e.speculative_runner, "_coda", original) + e.step() + assert seen[-1].owners[0] is not old.owners[0] + assert e.speculative_runner.stats.accepted_tokens >= accepted + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_failed_attention_releases_lease_and_allows_fresh_request(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + original = e.cache_manager._attend_prepared + + def fail(layer, kv, q): + if getattr(kv, "reference_plan", None) is not None: + raise RuntimeError("injected attention failure") + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", fail) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + with pytest.raises(RuntimeError, match="injected"): + for _ in range(100): + e.step() + assert e.reference_store.outstanding == 0 + assert "r" not in e.scheduler.requests + assert e.cache_manager.num_used_blocks == 0 + monkeypatch.setattr(e.cache_manager, "_attend_prepared", original) + generate(e, count=1) + finally: + e.close() + + +@pytest.mark.gpu +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cuda_sparse_sync_executes_and_drains(draft): + e = engine(draft=draft, retention=0.5, device="cuda", dtype=torch.bfloat16) + try: + generate(e) + torch.cuda.synchronize() + assert not bool(e.reference_store.counts.any()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_ordinary_retirement_survives_boundary_allocator_reuse(monkeypatch): + e = engine(device="cuda", dtype=torch.bfloat16) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + store = e.reference_store + assert store.outstanding == 1 + slot = e.model_runner.attention_references["r"][2].slots[0] + assert slot != 0 + torch.cuda.synchronize() + store.counts[0] = 77 + torch.cuda.synchronize() + clear = store._clear + pointers = [] + + def record(indices): + assert torch.cuda.current_stream() == e.model_runner.core_stream + pointers.append(indices.data_ptr()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + # Hold core while boundary allocations try to reuse any released index + # storage. Correct stream ownership prevents that reuse before clearing. + with torch.cuda.stream(e.model_runner.core_stream): + torch.cuda._sleep(20_000_000) + with torch.cuda.stream(e.model_runner.boundary_stream): + e.model_runner._release_attention_reference("r") + poison = [torch.zeros(1, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert pointers and len(poison) == 32 + assert not bool(store.counts[slot].any()) + assert bool((store.counts[0] == 77).all()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_close_drains_retirement_of_an_active_reference(): + e = engine(device="cuda", dtype=torch.bfloat16) + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + assert e.reference_store.outstanding == 1 + e.close() + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + assert e.model_runner.core_stream.query() diff --git a/vllm_rlt/__init__.py b/vllm_rlt/__init__.py index d0b346bc..c813eec0 100644 --- a/vllm_rlt/__init__.py +++ b/vllm_rlt/__init__.py @@ -5,6 +5,7 @@ ExecutionConfig, ExitConfig, SchedulerConfig, + SparseDecodeAttentionConfig, SpeculativeConfig, ) from vllm_rlt.entrypoints.llm import LLM @@ -19,6 +20,7 @@ "ExecutionConfig", "ExitConfig", "SchedulerConfig", + "SparseDecodeAttentionConfig", "SpeculativeConfig", "SamplingParams", "RequestOutput", diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..045c1a23 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -23,6 +23,8 @@ class BackendCapabilities: packed_matches_decode: bool = False # Consume GPU-owned context lengths without a host read between submissions. device_context_lengths: bool = False + # Shared paged sparse-reference capture/correction dispatch is available. + sparse_decode: bool = False @dataclass @@ -45,9 +47,9 @@ def generation(self): return getattr(self.kernel, "generation", None) -_TORCH_CAPABILITIES = BackendCapabilities() +_TORCH_CAPABILITIES = BackendCapabilities(sparse_decode=True) _TRITON_CAPABILITIES = BackendCapabilities( - cuda_graphs=True, async_scheduling=True, device_context_lengths=True + cuda_graphs=True, async_scheduling=True, device_context_lengths=True, sparse_decode=True ) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True), @@ -128,6 +130,7 @@ def create_backend( "cuda_graphs", "async_scheduling", "device_context_lengths", + "sparse_decode", ) if getattr(required_capabilities, name) and not getattr(implementation.capabilities, name) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py new file mode 100644 index 00000000..60c6abba --- /dev/null +++ b/vllm_rlt/attention/sparse.py @@ -0,0 +1,664 @@ +"""Sparse decode attention arithmetic and its per-row plan (RFC #156, section 3.2). + +The cross-loop sparse attention mechanism is from FlashLoop +(https://arxiv.org/abs/2609.29812, reference implementation +https://github.com/Superone77/FlashLoop); the kernels and serving integration in +this repository are ours. + +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. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + # Type-only: the attention layer borrows the runner-owned store, so it must + # not import the worker package at runtime. + from vllm_rlt.worker.attention_reference import ( + ReferenceLease, + SparseDecodeReferenceStore, + ) + +from vllm_rlt.kernels.sparse_capture import TRITON_AVAILABLE, fused_capture_scores +from vllm_rlt.kernels.sparse_consume import selected_scores as fused_selected_scores +from vllm_rlt.kernels.sparse_consume import selected_values as fused_selected_values + +# Row modes for one recurrent submission. +DENSE = 0 +CAPTURE = 1 +CONSUME = 2 + +# The reference loop is loop 2 in the RFC's one-based numbering, which is +# zero-based depth index 1 at runtime. +REFERENCE_DEPTH = 1 + + +def reference_mode(depth: int) -> int: + """Row mode of a recurrent step at the given zero-based loop depth.""" + if depth == REFERENCE_DEPTH: + return CAPTURE + if depth > REFERENCE_DEPTH: + return CONSUME + return DENSE + + +@dataclass(frozen=True) +class SparseDecodePlan: + """The reference plan of one recurrent submission. + + ``modes`` and ``slots`` are host-known: the runner decides both before it + submits, so a submission never depends on a device readback. ``slots`` holds + one reference slot per row, or ``-1`` for a row that owns no reference. + ``owners`` holds the lease each row's slot was reserved from, or ``None`` + for a dense row. The lease is the slot's identity: a slot handed to another + request is bound to a different lease, so a plan built for the previous + owner fails the store's liveness check. + + 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, ...] + owners: tuple[ReferenceLease | None, ...] + key_retention: float + slot_ids: torch.Tensor | None = None + store: SparseDecodeReferenceStore | 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} + + 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 slot_span(slots: Sequence[int]) -> slice | None: + """The slice covering ``slots`` when they are consecutive, else ``None``. + + Slot numbers follow the free-list order, not the row order, so a chunk of + rows can arrive fragmented (``[5, 3, 7]``). Comparing the endpoints would + accept that chunk — ``5 + 3 == 7 + 1`` — and then read or write ``[5, 6, 7]``, + which belongs to other requests. Every slot is checked against + ``first + offset`` instead. + """ + first = slots[0] + for offset, slot in enumerate(slots): + if slot != first + offset: + return None + return slice(first, first + len(slots)) + + +def expand_kv(tensor: torch.Tensor, groups: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention.""" + if groups == 1: + return tensor + return tensor.repeat_interleave(groups, dim=1) + + +def select_keys(probabilities: torch.Tensor, count: int) -> torch.Tensor: + """Pick the retained key positions of every head. + + ``probabilities`` is ``[heads, length]``. Equal probabilities are broken by + the smaller absolute position, so the selection does not depend on the order + of the rows in a batch. + """ + if count < 1: + raise ValueError("a selection needs at least one key") + if count > probabilities.shape[1]: + raise ValueError("cannot retain more keys than the query has") + # A stable descending sort keeps equal probabilities in position order. + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + return order[:, :count].sort(dim=-1).values + + +def reference_statistics( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + *, + key_retention: float, + key_bucket: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Dense reference attention and the retained state of one row. + + ``q`` is ``[heads, dim]``, ``keys`` and ``values`` are + ``[length, kv_heads, dim]``. Returns ``(output, contribution, mass, + positions)`` where output and contribution are ``[heads, dim]``, mass is + ``[heads]`` and positions are the retained absolute key positions. + """ + heads, dim = q.shape + length = keys.shape[0] + groups = heads // keys.shape[1] + scale = 1.0 / dim**0.5 + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q.float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output = torch.einsum("hl,lhd->hd", probabilities, expanded_values) + count = retained_count(length, key_retention) + # Buffer capacity never truncates the algorithm: a count above it fails + # loudly instead of silently degrading retention. + if key_bucket is not None and count > key_bucket: + raise ValueError(f"reference capacity {key_bucket} cannot hold {count} retained keys") + positions = select_keys(probabilities, count) + gathered = probabilities.gather(1, positions) + # Index the per-head value rows: positions is [heads, count], so the head + # index broadcasts against it and selects each head's own key rows. + head_index = torch.arange(heads, device=q.device)[:, None] + selected_values = expanded_values[positions, head_index] + contribution = torch.einsum("hk,hkd->hd", gathered, selected_values) + mass = gathered.sum(dim=-1) + return output, contribution, mass, positions + + +def sparse_correction( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + mass: torch.Tensor, + output: torch.Tensor, + contribution: torch.Tensor, +) -> torch.Tensor: + """Correct the reference output with the retained keys of a later loop. + + ``keys`` and ``values`` are ``[heads, retained, dim]`` and already carry the + keys of the later loop. ``mass``, ``output`` and ``contribution`` are the + retained reference state. + """ + dim = q.shape[-1] + scale = 1.0 / dim**0.5 + scores = torch.einsum("hd,hkd->hk", q.float(), keys.float()) * scale + probabilities = torch.softmax(scores, dim=-1) + new_contribution = torch.einsum("hk,hkd->hd", probabilities, values.float()) + return output - contribution + mass.float().unsqueeze(-1) * new_contribution + + +def gather_retained( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + selected: torch.Tensor, + block_size: int, + groups: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather each query head's retained keys and values from the paged cache. + + ``selected`` holds absolute key positions per query head, so every head + reads its own retained set. Returns ``[heads, retained, dim]`` tensors. + """ + index = selected.long() + blocks = block_table[index // block_size] + offsets = index % block_size + kv_index = (torch.arange(index.shape[0], device=index.device) // groups)[:, None] + kv_index = kv_index.expand_as(index) + return key_cache[blocks, offsets, kv_index], value_cache[blocks, offsets, kv_index] + + +def reference_paged_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + plan: SparseDecodePlan, + *, + layer: int, + dense_attention: Callable | None = None, +) -> torch.Tensor: + """Reference sparse decode attention over a paged KV cache (RFC #156 3.2). + + This is the numerical oracle for the sparse decode column, not an optimized + kernel: it walks rows on the host, reads the same paged tensors the dense + path reads, and publishes or consumes the retained state through the plan's + store. Dense rows reproduce the ordinary attention path exactly. + + ``key_cache``/``value_cache`` are one layer of + ``[physical block, token, KV head, dimension]``; ``block_tables`` and + ``context_lengths`` describe the submission rows exactly as the dense + backend receives them. When supplied, ``dense_attention`` preserves that + backend's outputs for dense/capture rows and the full-retention control. + Omitting it forces the independent Torch arithmetic oracle, including the + full-retention correction; that path need not be bitwise equal to Triton. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + block_size = key_cache.shape[1] + heads = q.shape[1] + groups = heads // key_cache.shape[2] + scale = 1.0 / math.sqrt(q.shape[-1]) + output = torch.empty_like(q) + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. Only the rows needing dense output use this result. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + captured = consumed = False + for row, length in enumerate(context_lengths.tolist()): + if length == 0: + output[row].zero_() + continue + positions = torch.arange(length, device=q.device) + blocks = block_tables[row, positions // block_size] + offsets = positions % block_size + mode = plan.modes[row] + if mode == CONSUME: + selected, mass, reference, contribution = store.read_reference(plan, layer, row) + if dense_output is not None and plan.key_retention == 1.0: + output[row] = dense_output[row] + consumed = True + continue + keys, values = gather_retained( + key_cache, value_cache, block_tables[row], selected, block_size, groups + ) + output[row] = sparse_correction(q[row], keys, values, mass, reference, contribution).to( + q.dtype + ) + consumed = True + continue + keys = key_cache[blocks, offsets] + values = value_cache[blocks, offsets] + if mode == CAPTURE: + reference, contribution, mass, selected = reference_statistics( + q[row], + keys, + values, + key_retention=plan.key_retention, + key_bucket=store.key_bucket, + ) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence. Keep + # selector probabilities/contributions in the Torch FP32 oracle. + reference = dense_output[row].float() + output[row] = reference.to(q.dtype) + store.write_reference( + plan, + layer, + row, + selected=selected, + counts=torch.full((heads,), selected.shape[-1], dtype=torch.int32, device=q.device), + mass=mass, + output=reference, + contribution=contribution, + ) + captured = True + continue + # Dense row: the ordinary attention path. + if dense_output is not None: + output[row] = dense_output[row] + continue + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q[row].float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output[row] = torch.einsum("hl,lhd->hd", probabilities, expanded_values).to(q.dtype) + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output + + +def sparse_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, +): + """Use bounded row batches on CUDA, retaining the scalar numerical oracle. + + A store configured for the ``fused`` implementation routes uniform capture + or consume submissions to the fused kernels. Other CUDA submissions keep + the batched path, whose capture rows take the fused paged score pass when a + dense backend is available; ``batched_paged_attention`` without + ``fused_scores`` still reproduces the scalar oracle bit for bit. + """ + if not q.is_cuda: + return reference_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + if getattr(plan.store, "implementation", "batched") == "fused": + from vllm_rlt.kernels import sparse_decode + + if sparse_decode.supports(plan, q, dense_attention): + return sparse_decode.fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + return batched_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + fused_scores=TRITON_AVAILABLE and dense_attention is not None, + ) + + +def _rowwise_product(equation, left, right): + """Keep each row's cuBLAS problem shape equal to the scalar oracle. + + Folding rows into the head dimension can choose a different FP32 reduction + kernel. Its small error can cross a BF16 rounding boundary and change a + greedy token, so batch the gathers/selection but retain row-wise products. + """ + # Prepare the exact views einsum uses once for the group, instead of + # repeating its permutation/reshape bookkeeping for every row. + if equation == "hd,lhd->hl": + right = right.permute(0, 2, 3, 1) + elif equation == "hl,lhd->hd": + right = right.transpose(1, 2) + elif equation == "hd,hkd->hk": + right = right.transpose(-1, -2) + elif equation != "hk,hkd->hd": + raise ValueError(f"unsupported reference product {equation}") + return torch.stack([torch.bmm(a, b).squeeze(1) for a, b in zip(left.unsqueeze(-2), right)]) + + +def batched_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, + fused_scores: bool = False, +): + """Batch equal-length rows without padding or changing retained-key selection. + + Each head still uses the same FP32 matrix products, stable sort and softmax + as the scalar oracle. Grouping exact lengths avoids padded reductions and + bounds gather scratch independently of the reference pool's key capacity. + There is one count readback per layer, rather than one per consume row. + + ``dense_attention`` keeps the numerical contract of the scalar path: dense + and capture rows publish the selected backend's output, the actually + rounded recurrence input becomes ``O_ref``, and a full-retention consume row + returns that same backend output as the control. Selector probabilities and + the retained contribution stay in the Torch FP32 oracle. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + store._require_live(plan) + lengths = context_lengths.tolist() + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + consume_rows = [r for r, mode in enumerate(plan.modes) if mode == CONSUME and lengths[r]] + counts = {} + if consume_rows: + # The retained count is host-known from the capture. Reading it back + # from the device would sync every layer, which RFC #156 section 6.2 + # forbids. + counts = {row: store.host_count(plan.slots[row], layer) for row in consume_rows} + missing = [row for row, count in counts.items() if count is None] + if missing: + # The fused capture decides its count on the device (graph replays + # change lengths without the host), so a mixed submission that + # falls back here after it reads those counts back once. + slots = plan.slot_ids[missing] + counts.update(zip(missing, store.counts[slots, layer].amax(dim=-1).tolist())) + if any(count == 0 for count in counts.values()): + raise RuntimeError("a consume row has no captured reference") + grouped = {} + for row, length in enumerate(lengths): + if length: + mode = plan.modes[row] + extent = counts[row] if mode == CONSUME else length + grouped.setdefault((mode, extent), []).append(row) + output = torch.zeros_like(q) + heads, dim = q.shape[1:] + groups = heads // key_cache.shape[2] + block_size = key_cache.shape[1] + scale = 1.0 / math.sqrt(dim) + head_index = torch.arange(heads, device=q.device)[None, :, None] + kv_index = head_index // groups + captured = consumed = False + for (mode, extent), all_rows in grouped.items(): + # The fused path materialises only the score matrix, so its budget is + # the scores themselves; the expanded path budgets an FP32 K row. Both + # must allow the whole group at once, or a 32k row degenerates into one + # launch per row. + if fused_scores: + chunk = max(1, min(32, 512 * 1024 * 1024 // (heads * extent * 4))) + else: + chunk = max(1, min(32, 64 * 1024 * 1024 // (extent * heads * dim))) + for start in range(0, len(all_rows), chunk): + rows = all_rows[start : start + chunk] + # Contiguous rows come from a device-side arange; otherwise a pinned + # source keeps the copy from draining the stream. + if rows[0] + len(rows) == rows[-1] + 1: + row_ids = torch.arange(rows[0], rows[0] + len(rows), device=q.device) + else: + row_ids = torch.tensor(rows, dtype=torch.int64, pin_memory=True).to( + q.device, non_blocking=True + ) + if mode == DENSE and dense_output is not None: + output[row_ids] = dense_output[row_ids] + continue + if mode == CONSUME and dense_output is not None and plan.key_retention == 1.0: + output[row_ids] = dense_output[row_ids] + consumed = True + continue + slots = plan.slot_ids[row_ids] + chunk_slots = [plan.slots[row] for row in rows] + queries = q[row_ids].float() + tables = block_tables[row_ids] + batch_index = torch.arange(len(rows), device=q.device)[:, None, None] + if mode == CONSUME: + if dense_output is not None and plan.key_retention == 1.0: + # Full retention is the control: same backend as dense. + output[row_ids] = dense_output[row_ids] + consumed = True + continue + # Consecutive slots turn these gathers into views, which is the + # difference between one index kernel per tensor and none. The + # order is checked per slot, never from the endpoints. + span = slot_span(chunk_slots) + if span is not None: + selected = store.selected[span, layer, :, :extent].to(torch.int32) + ref_output = store.output[span, layer] + ref_contribution = store.contribution[span, layer] + ref_mass = store.mass[span, layer] + else: + columns = torch.arange(extent, device=q.device)[None, None, :] + selected = store.selected[slots[:, None, None], layer, head_index, columns] + ref_output = store.output[slots, layer] + ref_contribution = store.contribution[slots, layer] + ref_mass = store.mass[slots, layer] + if fused_scores: + # Walk the retained indices in place: no dense K/V gather + # and no per-row cuBLAS. + scores = ( + fused_selected_scores( + q[row_ids], + key_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + * scale + ) + probabilities = torch.softmax(scores, dim=-1) + new = fused_selected_values( + probabilities, + value_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + else: + selected = selected.long() + blocks = tables[batch_index, selected // block_size] + offsets = selected % block_size + keys = key_cache[blocks, offsets, kv_index].float() + values = value_cache[blocks, offsets, kv_index].float() + scores = _rowwise_product("hd,hkd->hk", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + new = _rowwise_product("hk,hkd->hd", probabilities, values) + corrected = ref_output - ref_contribution + ref_mass.unsqueeze(-1) * new + output[row_ids] = corrected.to(q.dtype) + consumed = True + continue + positions = torch.arange(extent, device=q.device) + blocks = tables[:, positions // block_size] + offsets = positions % block_size + if fused_scores and dense_output is not None: + # One pass over the paged cache: no expanded FP32 K row, and + # no per-row cuBLAS. The dot product is still FP32; only the + # reduction order differs from the oracle. The context bound is + # host-known here, so the kernel never reads it back. + chunk_lengths = [lengths[row] for row in rows] + chunk_context = max(chunk_lengths) + scores = fused_capture_scores( + q[row_ids], + key_cache, + tables, + context_lengths[row_ids], + layer=layer, + block_size=block_size, + context=chunk_context, + ragged=min(chunk_lengths) != chunk_context, + ) + scores = scores * scale + else: + keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence; the + # selector state below stays in the Torch FP32 oracle. + reference = dense_output[row_ids].float() + else: + values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + reference = _rowwise_product("hl,lhd->hd", probabilities, values) + output[row_ids] = reference.to(q.dtype) + if mode == CAPTURE: + count = retained_count(extent, plan.key_retention) + store.ensure_key_capacity(count) + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + selected = order[..., :count].sort(dim=-1).values + gathered = probabilities.gather(-1, selected) + # Gather only retained V directly from the paged cache. The + # FP32 values and reduction geometry are identical to gathering + # from the full GQA-expanded tensor, without materializing it. + selected_blocks = tables[batch_index, selected // block_size] + selected_values = value_cache[ + selected_blocks, selected % block_size, kv_index + ].float() + contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) + # Only the live prefix is written, never the whole capacity. + span = slot_span(chunk_slots) + if span is not None: + store.selected[span, layer, :, :count] = selected.to(torch.int32) + store.counts[span, layer] = count + store.mass[span, layer] = gathered.sum(dim=-1) + store.output[span, layer] = reference + store.contribution[span, layer] = contribution + else: + columns = torch.arange(count, device=q.device)[None, None, :] + store.selected[slots[:, None, None], layer, head_index, columns] = selected.to( + torch.int32 + ) + store.counts[slots, layer] = count + store.mass[slots, layer] = gathered.sum(dim=-1) + store.output[slots, layer] = reference + store.contribution[slots, layer] = contribution + for row in rows: + store.note_count(plan.slots[row], layer, count) + captured = True + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 648781d7..521487ba 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 + # ``batched`` keeps the scalar oracle's reduction shapes (bitwise contract); + # ``fused`` launches fixed per-layer kernels without host readbacks and + # changes consume-row numerics (see vllm_rlt.kernels.sparse_decode). + implementation: str = "batched" + + def __post_init__(self): + if type(self.enabled) is not bool: + raise ValueError("enabled must be a boolean") + if self.implementation not in ("batched", "fused"): + raise ValueError("implementation must be 'batched' or 'fused'") + if not math.isfinite(self.key_retention) or not 0 < self.key_retention <= 1: + raise ValueError("key_retention must be in (0, 1]") + if self.key_bucket is not None: + _positive("key_bucket", self.key_bucket) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..a6595f28 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -23,6 +23,7 @@ create_backend, validate_backend_name, ) +from vllm_rlt.attention.sparse import sparse_paged_attention from vllm_rlt.request import ExecutionBatch from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder @@ -718,6 +719,21 @@ def _attend_prepared( self._validate_layer(layer) self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) + if not isinstance(batch, DeviceKVMetadata): + for allocation, depth, position in batch.rows: + self._require_prefix(allocation, layer, depth, position + 1) + plan = batch.reference_plan + if plan is not None and not plan.dense(): + return sparse_paged_attention( + q, + self.key_cache[:, layer], + self.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + dense_attention=self.attention, + ) if isinstance(batch, DeviceKVMetadata): return self.attention( q, @@ -728,8 +744,6 @@ def _attend_prepared( ) if not batch.rows: return torch.empty_like(q) - for allocation, depth, position in batch.rows: - self._require_prefix(allocation, layer, depth, position + 1) if batch.cu_seqlens_q is not None: return self.attention.prefill( q, diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..65ddc157 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 SparseDecodeReferenceStore from vllm_rlt.worker.executor import Executor from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput @@ -51,6 +55,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + sparse_decode_attention_config=None, stage_config=None, connector: StageConnector | None = None, ): @@ -87,6 +92,7 @@ def __init__( speculative_config and speculative_config.mode == "pipeline" ) self.async_speculative = bool(speculative_config and self.execution_config.async_scheduling) + scratch = 0 if speculative_config is not None: if cache_config.layout != "last_exited": raise ValueError("speculative decoding requires last_exited KV") @@ -118,6 +124,97 @@ 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_config = sparse_decode + self.reference_store = None + if sparse_decode is not None and sparse_decode.enabled: + if getattr(model, "requires_boundary_kv", False): + raise ValueError("sparse decode attention requires a recurrent Ouro-style model") + if cache_config.layout != "last_exited": + raise ValueError("sparse decode attention requires last_exited KV") + if self.execution_config.cuda_graphs: + raise ValueError("sparse decode attention does not yet support CUDA graphs") + if speculative_config is not None and speculative_config.mode != "block": + raise ValueError("sparse decode attention currently requires block speculation") + if scheduler_config.enable_preemption or scheduler_config.mode != "refill": + raise ValueError( + "sparse decode attention requires refill scheduling without preemption" + ) + if self.execution_config.async_scheduling and not self.async_speculative: + raise ValueError( + "sparse decode attention ordinary decode requires synchronous scheduling" + ) + # The index buffer never truncates the algorithm: capacity must + # hold the worst-case retained count over the model's positional + # limit. The buffer starts small and grows on demand, so a long + # model context does not cost its worst case up front. + context_bound = config.max_position_embeddings + required_keys = math.ceil(context_bound * sparse_decode.key_retention) + capacity = sparse_decode.key_bucket or required_keys + if capacity < required_keys: + raise ValueError( + f"sparse decode attention key_bucket {capacity} cannot hold " + f"{required_keys} retained keys for key_retention " + f"{sparse_decode.key_retention}; raise key_bucket or lower key_retention" + ) + # Pool capacity follows simultaneously live inputs, not requests + # alone. Synchronous block references retire in the same round; + # asynchronous speculation can retain two rounds per request. + slots = scheduler_config.max_num_seqs + if speculative_config is not None: + slots *= speculative_config.num_speculative_tokens + 1 + if self.async_speculative: + slots *= 2 + # Account for the reference pool before admitting any work: it is + # allocated up front, so a budget that cannot fit must fail here + # with an actionable message instead of OOM-ing mid-batch. + reference_bytes = ( + slots + * 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"key_retention or num_speculative_tokens" + ) + self.reference_store = SparseDecodeReferenceStore( + slots=slots, + layers=config.num_hidden_layers, + heads=config.num_attention_heads, + head_dim=config.head_dim, + key_bucket=capacity, + initial_key_capacity=min(capacity, 64), + device=parameter.device, + implementation=sparse_decode.implementation, + ) + if parameter.device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(parameter.device) + growth_bytes = self.reference_store.peak_capacity_bytes - self.reference_store.bytes + if growth_bytes > free_bytes: + raise ValueError( + "sparse reference index growth exceeds available device memory" + ) + logger.info( + "sparse decode attention: %d reference slots x %d keys x %d layers x %d heads " + "<= %.2f GiB growth peak (context bound %d, retention %.2f)", + slots, + capacity, + config.num_hidden_layers, + config.num_attention_heads, + self.reference_store.peak_capacity_bytes / 1024**3, + context_bound, + sparse_decode.key_retention, + ) + # Budget the logical maximum and copy overlap, not just the small + # initial allocation; short prompts must not overcommit long ones. + scratch += self.reference_store.peak_capacity_bytes + if speculative_config is not None or self.reference_store is not None: cache_config = replace( cache_config, memory_reserve_bytes=cache_config.memory_reserve_bytes + scratch ) @@ -167,6 +264,7 @@ def __init__( parameter.device.type == "cuda" and self.execution_config.async_scheduling ), device_context_lengths=self.async_speculative, + sparse_decode=self.reference_store is not None, ), ) num_blocks, self.memory_plan = plan_cache( @@ -217,6 +315,8 @@ def __init__( scheduler_config=scheduler_config, staged_prefill=self.scheduler.staged_prefill, rng=self.rng, + reference_store=self.reference_store if speculative_config is None else None, + sparse_decode_config=self.sparse_decode_config, ) self.speculative_runner = ( SpeculativeRunner( @@ -227,6 +327,8 @@ def __init__( rng=self.rng, runtime=self.model_runner.runtime, compute_stream=self.model_runner.core_stream, + reference_store=self.reference_store, + sparse_decode_config=self.sparse_decode_config, ) if speculative_config is not None and not self.async_speculative else None @@ -244,7 +346,15 @@ def __init__( copy_stream=self.model_runner.copy_stream, rng=self.rng, runtime=self.model_runner.runtime, + reference_store=self.reference_store, + sparse_decode_config=self.sparse_decode_config, ) + if self.reference_store is not None: + # Reference lifecycle device work is ordered on the compute + # stream, after the readers that retire a slot. + self.reference_store.stream = self.model_runner.core_stream + if self.reference_store is not None: + self.reference_store.stream = self.model_runner.core_stream self.executor = Executor(self.model_runner, speculation=self.speculative_runner) self._pending_speculative = [] self._retiring_speculative = {} @@ -315,8 +425,12 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i ) and params.temperature != 0: mode = "async" if self.async_speculative else "pipeline" raise ValueError(f"{mode} speculative decoding currently requires greedy sampling") - if request_id in self._retiring_speculative: - self._retire_speculative(wait=True, request_id=request_id) + # A request id is never reused while its previous token is retiring: + # reference slots are recycled only after their last reader completed, + # so rejecting the duplicate keeps ownership unambiguous without a + # per-slot generation counter. + if request_id in self._retiring_speculative or request_id in self.scheduler.requests: + raise ValueError(f"duplicate request ID {request_id!r}: already live or still retiring") config = self.model.config if not prompt_token_ids: raise ValueError("prompt must contain at least one token") @@ -400,6 +514,13 @@ def close(self): self.abort_request(rid) if self.async_speculative: self._discard_speculative() + if self.reference_store is not None: + # Abort can enqueue reference metadata retirement after the drain + # above. Teardown must also cover those final writes. + self.model_runner.synchronize() + # The store owns the stream the fused kernels ran on, so it releases + # their module-level scratch here — after the drain, never before. + self.reference_store.close() def _discard_speculative(self): """Release speculative owners after the caller has drained both streams.""" 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..58ae437c 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,17 @@ 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( "--cuda-graphs", action="store_true", @@ -129,6 +146,12 @@ 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, + ) + if args.sparse_decode_attention + else None, ) diff --git a/vllm_rlt/kernels/sparse_capture.py b/vllm_rlt/kernels/sparse_capture.py new file mode 100644 index 00000000..e9671f35 --- /dev/null +++ b/vllm_rlt/kernels/sparse_capture.py @@ -0,0 +1,152 @@ +"""Fused paged score pass for the sparse-decode capture step. + +The selector needs a probability for every key of the reference loop, but the +Torch oracle materialises the whole causal row in FP32 with the GQA heads +expanded, which costs tens of MiB per layer and leaves the pass CPU-bound. This +kernel reads the paged cache once and writes one FP32 score per +``(row, head, position)``, so selection, softmax and the retained contribution +stay on the device without an expanded copy. + +It computes the same dot product as the oracle in FP32 (the cache is BF16, and +BF16 -> FP32 is exact); only the reduction order differs. Measured against the +oracle at 8k context: max score delta 1.1e-5 on a magnitude of 65, retained mass +delta 9e-8, full attention output delta 2e-8, and identical selected sets on +511/512 row-head pairs. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _capture_scores_kernel( + q_ptr, + k_ptr, + t_ptr, + l_ptr, + out_ptr, + context, + q_stride_row, + k_stride_block, + DIM: tl.constexpr, + GROUPS: tl.constexpr, + KV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + row = tl.program_id(0) + kv_head = tl.program_id(1) + chunk = tl.program_id(2) + + position = chunk * BLOCK_N + tl.arange(0, BLOCK_N) + row_length = tl.load(l_ptr + row) + mask = position < row_length + page = tl.load(t_ptr + row * PAGES + position // BLOCK_SIZE, mask=mask, other=0) + offset = position % BLOCK_SIZE + dim = tl.arange(0, DIM) + # [BLOCK_N, DIM] + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + for group in tl.static_range(GROUPS): + head = kv_head * GROUPS + group + query = tl.load(q_ptr + row * q_stride_row + head * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store( + out_ptr + (row * (KV * GROUPS) + head) * context + position, + scores, + mask=mask, + ) + + +def fused_capture_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + *, + layer: int, + block_size: int, + block_n: int = 128, + context: int | None = None, + ragged: bool | None = None, +) -> torch.Tensor: + """Scores of every row against its causal keys, straight from the cache. + + ``queries`` is ``[rows, heads, dim]``, ``key_cache`` is either + ``[blocks, layers, block_size, kv_heads, dim]`` or the already sliced + ``[blocks, block_size, kv_heads, dim]``, and ``block_tables`` is + ``[rows, pages]``. Returns ``[rows, heads, max_context]`` FP32 with entries + beyond each row's length left at zero. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused capture needs Triton") + if not queries.is_cuda: + raise ValueError("fused capture is a CUDA path") + + rows, heads, dim = queries.shape + if key_cache.dim() == 5: + layer_cache = key_cache[:, layer] + kv_heads = key_cache.shape[3] + else: + # The cache manager hands the attention layer a per-layer view already. + layer_cache = key_cache + kv_heads = key_cache.shape[2] + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + lengths = context_lengths.to(torch.int64) + # ``context`` and ``ragged`` are host-known for a runner submission. Reading + # them back with ``.item()`` would sync the stream once per layer, which + # RFC #156 section 6.2 forbids, so callers pass them in. + if context is None: + context = int(lengths.max().item()) if lengths.numel() else 0 + if context == 0: + return queries.new_zeros((rows, heads, 0), dtype=torch.float32) + + pages = block_tables.shape[1] + out = torch.zeros(rows, heads, context, device=queries.device, dtype=torch.float32) + grid = (rows, kv_heads, triton.cdiv(context, block_n)) + _capture_scores_kernel[grid]( + queries.contiguous(), + layer_cache, + block_tables, + lengths, + out, + context, + queries.stride(0), + key_cache.stride(0), + DIM=dim, + GROUPS=groups, + KV=kv_heads, + BLOCK_SIZE=block_size, + PAGES=pages, + BLOCK_N=block_n, + ) + # Rows are ragged; the kernel leaves the tail at zero, so mask it to -inf + # before any softmax sees it. + if ragged is None: + ragged = lengths.numel() > 0 and int(lengths.min().item()) != context + if ragged: + positions = torch.arange(context, device=queries.device)[None, None, :] + out = out.masked_fill(positions >= lengths[:, None, None], float("-inf")) + return out diff --git a/vllm_rlt/kernels/sparse_consume.py b/vllm_rlt/kernels/sparse_consume.py new file mode 100644 index 00000000..06b7fc7b --- /dev/null +++ b/vllm_rlt/kernels/sparse_consume.py @@ -0,0 +1,209 @@ +"""Fused gather-and-accumulate for the sparse-decode consume step. + +A consume row only touches its ``k`` retained keys, but the Torch path gathers +them into a dense ``[rows, heads, k, dim]`` tensor twice (K then V) and runs one +cuBLAS call per row. At 32k context with ``k = 3277`` that is 8.2 ms of K gather +plus 8.1 ms of V gather through three advanced indices on a strided +``[blocks, block_size, kv, dim]`` cache, which measures about 45 GB/s. + +These two kernels walk the retained indices in place: the first writes one FP32 +score per ``(row, head, retained)``, the second accumulates ``p * V`` into one +output row. Nothing dense is materialised, and the softmax in between stays in +Torch because ``k`` is small. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _selected_scores_kernel( + q_ptr, + k_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + k_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + block = tl.program_id(2) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + + slot = block * BLOCK_K + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + dim = tl.arange(0, DIM) + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + query = tl.load(q_ptr + (row * heads + head) * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store(out_ptr + (row * heads + head) * retained + slot, scores, mask=mask) + + @triton.jit + def _selected_values_kernel( + p_ptr, + v_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + v_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + dim = tl.arange(0, DIM) + accumulator = tl.zeros([DIM], dtype=tl.float32) + for start in range(0, retained, BLOCK_K): + slot = start + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + weight = tl.load( + p_ptr + (row * heads + head) * retained + slot, mask=mask, other=0.0 + ).to(tl.float32) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + values = tl.load( + v_ptr + + page[:, None] * v_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + accumulator += tl.sum(weight[:, None] * values, axis=0) + tl.store(out_ptr + (row * heads + head) * DIM + dim, accumulator) + + +def _cache_view(cache: torch.Tensor, layer: int) -> tuple[torch.Tensor, int]: + """Return the per-layer cache view and its KV head count.""" + if cache.dim() == 5: + return cache[:, layer], cache.shape[3] + return cache, cache.shape[2] + + +def selected_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``q . K[S]`` for every retained key, without materialising K. + + ``queries`` is ``[rows, heads, dim]``, ``selected`` is ``[rows, heads, k]`` + of absolute positions and the result is ``[rows, heads, k]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, dim = queries.shape + cache, kv_heads = _cache_view(key_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + retained = selected.shape[-1] + out = torch.empty(rows, heads, retained, device=queries.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads, triton.cdiv(retained, block_k)) + _selected_scores_kernel[grid]( + queries.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out + + +def selected_values( + probabilities: torch.Tensor, + value_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``sum_j p[j] * V[S[j]]`` for every row and head. + + ``probabilities`` is ``[rows, heads, k]`` and the result is + ``[rows, heads, dim]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, retained = probabilities.shape + cache, kv_heads = _cache_view(value_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + dim = cache.shape[-1] + out = torch.empty(rows, heads, dim, device=probabilities.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads) + _selected_values_kernel[grid]( + probabilities.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py new file mode 100644 index 00000000..d0c8ffd0 --- /dev/null +++ b/vllm_rlt/kernels/sparse_decode.py @@ -0,0 +1,1385 @@ +"""Fused sparse decode attention for uniform capture or consume submissions. + +Opt-in alternative to the batched path in ``vllm_rlt.attention.sparse``. That +path keeps the scalar oracle's cuBLAS reduction shapes, so every layer launches +per-row products and reads row lengths and retained counts back to the host. +Here each layer launches a fixed number of kernels and keeps lengths and +counts on the device. + +Numerics differ from the oracle on consume rows: scores and FP32 reductions +run in another order, so a key on the selection boundary can differ from the +oracle's. Retained positions are written in ascending order and equal scores +keep the smaller positions, as in the oracle. Capture rows still return the +dense backend output, which is also published as ``O_ref``, so the recurrence +input of the reference loop is unchanged. + +A consume row whose reference was never captured (count 0) yields NaN instead +of raising: raising would need the host readback this path exists to avoid. The +batched and scalar paths keep the explicit error. +""" + +import operator + +import torch +import triton +import triton.language as tl + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, retained_count +from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention +from vllm_rlt.kernels.sparse_workspace import CAPTURE_SCRATCH, CONSUME_SCRATCH, RETENTION_SCALARS +from vllm_rlt.kernels.triton_attention import NUM_SPLITS + + +@triton.jit +def _order_key(scores): + """Key in [0, 2**32) with the order of FP32 scores (bits of negatives flipped).""" + bits = scores.to(tl.int32, bitcast=True) + return (bits ^ ((bits >> 31) & 0x7FFFFFFF)).to(tl.int64) + 2147483648 + + +@triton.jit +def _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + reference, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Keep the ``count`` largest of ``length`` stored scores and publish ``(S, m, O_ref, C_ref)``. + + Finds the count-th largest score by a radix select over its 32-bit order + key, four bits per pass, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. ``reference`` is ``O_ref``. + """ + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + # After eight passes ``prefix`` is the order key of the count-th largest + # score and ``ties_needed`` how many keys equal to it are still retained. + # Loop-carried values must keep one type on every iteration. + prefix = length.to(tl.int64) * 0 + ties_needed = count + digits = tl.arange(0, 16) + for level in tl.static_range(8): + histogram = tl.zeros((16,), dtype=tl.int32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 4 * level)) == prefix) + digit = ((key >> (28 - 4 * level)) & 15).to(tl.int32) + histogram += tl.sum( + ((digit[:, None] == digits[None, :]) & candidate[:, None]).to(tl.int32), axis=0 + ) + # Highest digit whose keys, with every higher digit, still cover the rest. + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + prefix = prefix * 16 + chosen + low = prefix + selected = SELECTED + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head) * retained_capacity + written = 0 + ties_seen = 0 + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + scores = tl.load(scores_row + positions, mask=live, other=0) + key = _order_key(scores) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + probabilities = tl.where(take, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=take, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=take[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + tl.store(MASS + state, mass) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, contribution, mask=dims < HEAD_DIM) + + +# A copy of triton_attention._paged_attention_kernel that also stores the raw +# scores. Its arithmetic must stay identical to the dense kernel: the capture +# output is required to equal the dense backend bit for bit (tested). +@triton.jit +def _capture_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + PART_M, + PART_L, + PART_O, + SCORES, + SEGHIST, + RETENTION, + STATE, + TICKETS, + top, + score_capacity, + q_batch_stride: tl.constexpr, + q_head_stride: tl.constexpr, + q_dim_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + k_dim_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + v_dim_stride: tl.constexpr, + table_stride: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + NUM_HEADS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, + HISTOGRAM: tl.constexpr, + PARTS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + kv_head = head // HEAD_GROUPS + length = tl.load(LENGTHS + row) + # Equal tile-aligned parts of this row's context; trailing parts may be empty. + split_tokens = ((length + SPLITS - 1) // SPLITS + BLOCK_T - 1) // BLOCK_T * BLOCK_T + begin = split * split_tokens + end = tl.minimum(length, begin + split_tokens) + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_batch_stride + head * q_head_stride + dims * q_dim_stride, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + local = tl.zeros((512,), dtype=tl.int32) + for start in range(begin, end, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + valid_tokens = positions < end + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, + mask=valid_tokens, + other=0, + ).to(tl.int64) + # Physical IDs fit in int32, but ID * block_stride can exceed 2**31 + # elements in a large KV pool. Promote BEFORE multiplying, for K and V. + offsets = positions % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :] * k_dim_stride, + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) * SCALE + scores = tl.where(valid_tokens, scores, -float("inf")) + # The only addition to the dense kernel: keep the raw scores for selection. + tl.store( + SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + positions, + scores, + mask=valid_tokens, + ) + if HISTOGRAM: + # Selection pass 0 of this part: the top eight bits of the order key. + digit = (_order_key(scores) >> 24).to(tl.int32) + local += tl.histogram(tl.where(valid_tokens, digit, 256), 512) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :] * v_dim_stride, + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) + normalizer = normalizer * correction + tl.sum(weights, axis=0) + maximum = next_maximum + # An empty part stores (-inf, 0, 0), which the merge weights by zero. + part = (row * NUM_HEADS + head) * SPLITS + split + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + if HISTOGRAM: + rh = row * NUM_HEADS + head + _, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + bins = tl.arange(0, 512) + tl.store(SEGHIST + (rh * PARTS + split) * 4 * 256 + bins, local, mask=bins < 256) + _close_pass(SEGHIST, STATE, TICKETS, rh, count, 0, SPLITS, PARTS) + + +# Parallel selection after the split kernel. Every (row, head) splits its +# context into PARTS tile-aligned parts. Four radix passes over the 32-bit order +# key, eight bits each, store one 256-bin histogram per part. The last part to +# finish a pass (an integer ticket per row and head) sums the histograms in a +# fixed order and publishes the state for the next pass, so the result stays +# deterministic and no histogram is zeroed. After the last pass it also derives +# every part's output offset and tie quota, which keeps the ascending order and +# the smaller-position tie rule. The contribution kernel then gathers V over the +# compact retained list only. +SELECT_PARTS = 16 +# Tiles of the radix passes and the write (scores only) and of the V gather, +# chosen by a sweep on an RTX 5090 at 4K-32K (16 vs 32 parts, 128-1024 scan, +# 32-128 gather, 4 or 8 warps). +_SCAN_BLOCK = 1024 +_SCAN_WARPS = 4 +_GATHER_BLOCK = 128 +# The split kernel's tile (the dense kernel's), so its splits are whole parts. +_SEGMENT_ALIGN = 32 +# The split kernel builds the first pass's histograms while it scores the keys. +_PASS0_IN_SPLIT = True + +# Triton's Python dispatch costs several microseconds per launch, and a fused +# capture makes six launches, so at short contexts in eager mode the dispatch, +# not the GPU, sets its cost. A launch whose specialization was seen before +# reuses that compiled kernel directly. The key follows Triton's own +# specialization (dtype and 16-byte alignment of pointers; integers equal to 1, +# divisible by 16, and their width; constexpr values); tests set _LAUNCH_CHECK +# to assert that it always picks the kernel Triton would. Other Triton versions +# keep its dispatch. +try: + from triton import knobs as _knobs + from triton.runtime.driver import driver as _driver + + _FAST_LAUNCH = triton.__version__.startswith("3.8.") and hasattr( + _knobs.runtime, "launch_enter_hook" + ) +except ImportError: # pragma: no cover - other Triton versions + _FAST_LAUNCH = False +_LAUNCH_CHECK = False + + +def _key_function(args, constexpr): + """Compile the specialization key of one argument layout, as Triton's binder does. + + The layout is the arguments' types; each layout gets its own function. + """ + terms = [] + for i, value in enumerate(args): + if i in constexpr: + terms.append(f"a[{i}]") + elif isinstance(value, torch.Tensor): + terms.append(f"a[{i}].dtype, a[{i}].data_ptr() & 15 == 0") + elif type(value) is int: + terms.append(f"a[{i}] == 1, a[{i}] & 15 == 0, -2147483648 <= a[{i}] < 2147483648") + elif type(value) is not float: + return None + if not terms: + return lambda a: () + body = ", ".join(terms) + return eval(f"lambda a: ({body},)") # noqa: S307 - built from indices only + + +class _Launcher: + """``kernel[grid](*args, **meta)`` that skips Triton's dispatch once compiled.""" + + def __init__(self, kernel): + self.kernel = kernel + self.compiled = {} + self.keys = {} + # Positional parameters declared constexpr specialize on their value. + self.constexpr = frozenset(i for i, p in enumerate(kernel.params) if p.is_constexpr) + + def launch(self, grid, args, meta): + if not _FAST_LAUNCH: + self.kernel[grid](*args, **meta) + return + layout = (tuple(map(type, args)), len(grid)) + keyed = self.keys.get(layout) + if keyed is None: + keyed = self.keys[layout] = _key_function(args, self.constexpr) + if keyed is None: + self.kernel[grid](*args, **meta) + return + device = _driver.active.get_current_device() + key = (device, layout, keyed(args), tuple(meta.items())) + entry = self.compiled.get(key) + if entry is None: + compiled = self.kernel[grid](*args, **meta) + # The launcher takes the bound arguments in the binder's order. + binder = self.kernel.device_caches[device][4] + options = dict( + meta, + debug=self.kernel.debug or _knobs.runtime.debug, + instrumentation_mode=_knobs.compilation.instrumentation_mode, + ) + names = tuple(binder(*args, **options)[0])[len(args) :] + # None when the call passes its meta in the binder's order. + if names == tuple(meta): + order = None + elif len(names) > 1: + order = operator.itemgetter(*names) + else: + order = lambda meta: tuple(meta[name] for name in names) # noqa: E731 + self.compiled[key] = (compiled, order) + return + compiled, order = entry + if _LAUNCH_CHECK: + assert self.kernel.run(*args, grid=grid, warmup=True, **meta) is compiled + values = args + (tuple(meta.values()) if order is None else order(meta)) + stream = _driver.active.get_current_stream(device) + compiled.run( + grid[0], + grid[1] if len(grid) > 1 else 1, + grid[2] if len(grid) > 2 else 1, + stream, + compiled.function, + compiled.packed_metadata, + compiled.launch_metadata(grid, stream, *values), + _knobs.runtime.launch_enter_hook, + _knobs.runtime.launch_exit_hook, + *values, + ) + + +def _current_stream(device): + """The raw current stream; Triton's driver query is much cheaper than torch's.""" + if _FAST_LAUNCH: + return _driver.active.get_current_stream(device.index) + return torch.cuda.current_stream(device).cuda_stream + + +_LAUNCHERS: dict = {} + + +def _launch(kernel, grid, *args, **meta): + # Keyed by identity: hashing a JITFunction recomputes its cache key. + launcher = _LAUNCHERS.get(id(kernel)) + if launcher is None or launcher.kernel is not kernel: + launcher = _LAUNCHERS[id(kernel)] = _Launcher(kernel) + launcher.launch(grid, args, meta) + + +@triton.jit +def _retained(LENGTHS, row, RETENTION, top): + """The row length and its retained count, as the host's ``retained_count``.""" + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + return length, tl.minimum(count.to(tl.int32), top) + + +@triton.jit +def _segment(length, SEGMENTS: tl.constexpr, ALIGN: tl.constexpr): + """This program's part ``[first, last)`` of a row's context.""" + size = tl.cdiv(tl.cdiv(length, SEGMENTS), ALIGN) * ALIGN + first = tl.program_id(2) * size + return first, tl.minimum(first + size, length) + + +@triton.jit +def _choose(histogram, prefix, ties_needed): + """Highest digit whose keys, with every higher digit, still cover the rest.""" + digits = tl.arange(0, 256) + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + return prefix * 256 + chosen.to(tl.int64), ties_needed, chosen + + +@triton.jit +def _last_part(TICKETS, rh, PARTS: tl.constexpr): + """Whether this program is the last of its row and head to finish; resets the ticket. + + Each part stores its results first; the acquire-release ticket then makes + every part's stores visible to the last one, which reads them through L2. + """ + tl.debug_barrier() + last = tl.atomic_add(TICKETS + rh, 1, sem="acq_rel", scope="gpu") == PARTS - 1 + if last: + tl.store(TICKETS + rh, 0) + return last + + +@triton.jit +def _close_pass( + SEGHIST, STATE, TICKETS, rh, count, PASS: tl.constexpr, USED: tl.constexpr, PARTS: tl.constexpr +): + """The last part of a pass publishes the state before the next pass.""" + if _last_part(TICKETS, rh, USED): + digits = tl.arange(0, 256) + parts = (rh * PARTS + tl.arange(0, USED)) * 4 + PASS + histogram = tl.sum( + tl.load(SEGHIST + parts[:, None] * 256 + digits[None, :], cache_modifier=".cg"), axis=0 + ) + if PASS == 0: + prefix = count.to(tl.int64) * 0 + ties_needed = count + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + ties_needed = tl.load(STATE + (rh * 5 + PASS) * 2 + 1).to(tl.int32) + prefix, ties_needed, _ = _choose(histogram, prefix, ties_needed) + tl.store(STATE + (rh * 5 + PASS + 1) * 2, prefix) + tl.store(STATE + (rh * 5 + PASS + 1) * 2 + 1, ties_needed.to(tl.int64)) + + +@triton.jit +def _select_pass_kernel( + PART_M, + PART_L, + PART_O, + OUT, + SCORES, + LENGTHS, + RETENTION, + SLOTS, + COUNTS, + OUTPUT, + LSE, + SEGHIST, + STATE, + TICKETS, + GREATER, + OFFSETS, + layer, + top, + score_capacity, + out_batch_stride: tl.constexpr, + out_head_stride: tl.constexpr, + out_dim_stride: tl.constexpr, + PASS: tl.constexpr, + MERGE: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, + SPLITS: tl.constexpr, +): + """One radix pass of one part; the first launched pass also merges the split parts.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if MERGE: + if segment == 0: + # Identical to triton_attention._paged_attention_merge_kernel. + parts = rh * SPLITS + tl.arange(0, SPLITS) + dims = tl.arange(0, BLOCK_D) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + dense = accumulator / tl.maximum(normalizer, 1.0e-30) + out = OUT + row * out_batch_stride + head * out_head_stride + dims * out_dim_stride + tl.store(out, dense, mask=dims < HEAD_DIM) + slot = tl.load(SLOTS + row) + tl.store(LSE + rh, maximum + tl.log(normalizer)) + if count > 0: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + # O_ref is the output as rounded to the query dtype: the recurrence input. + reference = dense.to(OUT.dtype.element_ty).to(tl.float32) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + if count > 0: + if PASS == 0: + prefix = length.to(tl.int64) * 0 + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + first, last = _segment(length, PARTS, ALIGN) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + local = tl.zeros((512,), dtype=tl.int32) + greater = 0 + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 8 * PASS)) == prefix) + digit = ((key >> (24 - 8 * PASS)) & 255).to(tl.int32) + # Bin 256 collects the non-candidates and is never stored. + local += tl.histogram(tl.where(candidate, digit, 256), 512) + if PASS == 3: + # Keys of this part already above the threshold's first 24 bits. + greater += tl.sum((live & ((key >> 8) > prefix)).to(tl.int32), axis=0) + bins = tl.arange(0, 512) + own = (rh * PARTS + segment) * 4 + PASS + tl.store(SEGHIST + own * 256 + bins, local, mask=bins < 256) + if PASS == 3: + tl.store(GREATER + rh * PARTS + segment, greater) + if _last_part(TICKETS, rh, PARTS): + # Fix the threshold, then every part's offset and tie quota. + digits = tl.arange(0, 256) + segments = tl.arange(0, PARTS) + prefix = tl.load(STATE + (rh * 5 + 3) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 3) * 2 + 1).to(tl.int32) + last_pass = tl.load( + SEGHIST + ((rh * PARTS + segments[:, None]) * 4 + 3) * 256 + digits[None, :], + cache_modifier=".cg", + ) + low, ties_needed, chosen = _choose(tl.sum(last_pass, axis=0), prefix, ties_needed) + above = tl.load(GREATER + rh * PARTS + segments, cache_modifier=".cg") + tl.sum( + tl.where(digits[None, :] > chosen, last_pass, 0), axis=1 + ) + equal = tl.sum(tl.where(digits[None, :] == chosen, last_pass, 0), axis=1) + ties_before = tl.cumsum(equal, 0) - equal + taken = above + tl.minimum(tl.maximum(ties_needed - ties_before, 0), equal) + tl.store(OFFSETS + (rh * PARTS + segments) * 2, tl.cumsum(taken, 0) - taken) + tl.store(OFFSETS + (rh * PARTS + segments) * 2 + 1, ties_before) + tl.store(STATE + (rh * 5 + 4) * 2, low) + tl.store(STATE + (rh * 5 + 4) * 2 + 1, ties_needed.to(tl.int64)) + else: + _close_pass(SEGHIST, STATE, TICKETS, rh, count, PASS, PARTS, PARTS) + + +@triton.jit +def _select_write_kernel( + SCORES, + LENGTHS, + RETENTION, + SLOTS, + STATE, + OFFSETS, + SELECTED, + layer, + top, + score_capacity, + retained_capacity, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, +): + """Write this part's retained positions at its offset, ties up to its quota.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + low = tl.load(STATE + (rh * 5 + 4) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 4) * 2 + 1).to(tl.int32) + written = tl.load(OFFSETS + (rh * PARTS + segment) * 2) + ties_seen = tl.load(OFFSETS + (rh * PARTS + segment) * 2 + 1) + slot = tl.load(SLOTS + row) + selected = ( + SELECTED + + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head).to(tl.int64) * retained_capacity + ) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + first, last = _segment(length, PARTS, ALIGN) + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + + +@triton.jit +def _select_contribution_kernel( + SCORES, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + LSE, + SELECTED, + PART_MASS, + PART_C, + TICKETS, + MASS, + CONTRIBUTION, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, +): + """``(m, C_ref)`` over the retained ranks, in rank chunks summed in a fixed order. + + Gathering V over the compact retained list keeps every lane busy, where a + scan over all positions would leave most of them masked. + """ + row = tl.program_id(0) + head = tl.program_id(1) + part = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + chunk = tl.cdiv(tl.cdiv(count, PARTS), BLOCK_T) * BLOCK_T + first = part * chunk + last = tl.minimum(first + chunk, count) + lse = tl.load(LSE + rh) + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + selected = SELECTED + state.to(tl.int64) * retained_capacity + scores_row = SCORES + rh.to(tl.int64) * score_capacity + table = TABLES + row * table_stride + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + scores = tl.load(scores_row + positions, mask=live, other=0) + probabilities = tl.where(live, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + index = rh * PARTS + part + tl.store(PART_MASS + index, mass) + tl.store(PART_C + index * BLOCK_D + dims, contribution) + if _last_part(TICKETS, rh, PARTS): + parts = rh * PARTS + tl.arange(0, PARTS) + total = tl.sum(tl.load(PART_MASS + parts, cache_modifier=".cg"), axis=0) + summed = tl.sum( + tl.load(PART_C + parts[:, None] * BLOCK_D + dims[None, :], cache_modifier=".cg"), + axis=0, + ) + tl.store(MASS + state, total) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, summed, mask=dims < HEAD_DIM) + + +@triton.jit +def _capture_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + DENSE, + SCORES, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + dense_row_stride: tl.constexpr, + dense_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Select and publish the retained state ``(S, m, O_ref, C_ref)`` of one row and head. + + One program scores every key, finds the ``count``-th largest score by + bisection on its order key, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. + """ + row = tl.program_id(0) + head = tl.program_id(1) + slot = tl.load(SLOTS + row) + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + # key_bucket covers the serving context, so this never truncates; it only + # keeps a misconfiguration from writing past the index buffer. + count = tl.minimum(count.to(tl.int32), top) + if count > 0: + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + table = TABLES + row * table_stride + scores_row = SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + maximum = -float("inf") + normalizer = 0.0 + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, other=0).to(tl.int64) + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + (positions % PAGE_SIZE)[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + tl.store(scores_row + positions, scores, mask=live) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + normalizer = normalizer * tl.exp(maximum - next_maximum) + tl.sum( + tl.exp(scores - next_maximum), axis=0 + ) + maximum = next_maximum + lse = maximum + tl.log(normalizer) + # Later passes read scores other threads of this program stored. + tl.debug_barrier() + dense = tl.load( + DENSE + row * dense_row_stride + head * dense_head_stride + dims, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + dense, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride, + v_token_stride, + v_head_stride, + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + HEAD_GROUPS, + PAGE_SIZE, + BLOCK_D, + BLOCK_T, + ) + + +@triton.jit +def _consume_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + SLOTS, + SELECTED, + COUNTS, + PART_M, + PART_L, + PART_O, + layer, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, +): + """Online-softmax parts of ``C_new`` over one contiguous chunk of retained ranks. + + Each (row, head) splits its ``count`` retained keys into ``SPLITS`` chunks + of whole ``BLOCK_T`` tiles; an empty chunk writes the merge's neutral part. + """ + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + dims = tl.arange(0, BLOCK_D) + part = (row * NUM_HEADS + head) * SPLITS + split + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + if tl.load(LENGTHS + row) > 0: + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + count = tl.load(COUNTS + state) + chunk = tl.cdiv(tl.cdiv(count, SPLITS), BLOCK_T) * BLOCK_T + first = split * chunk + last = tl.minimum(first + chunk, count) + selected = SELECTED + state.to(tl.int64) * retained_capacity + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, mask=live, other=0 + ).to(tl.int64) + offsets = positions % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) + normalizer = normalizer * correction + tl.sum(weights, axis=0) + maximum = next_maximum + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + + +@triton.jit +def _consume_merge_kernel( + PART_M, + PART_L, + PART_O, + LENGTHS, + SLOTS, + MASS, + OUTPUT, + CONTRIBUTION, + OUT, + layer, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + SPLITS: tl.constexpr, +): + """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order.""" + row = tl.program_id(0) + head = tl.program_id(1) + dims = tl.arange(0, BLOCK_D) + out = OUT + row * q_row_stride + head * q_head_stride + dims + slot = tl.load(SLOTS + row) + if tl.load(LENGTHS + row) == 0: + tl.store(out, tl.zeros((BLOCK_D,), dtype=tl.float32), mask=dims < HEAD_DIM) + else: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + parts = (row * NUM_HEADS + head) * SPLITS + tl.arange(0, SPLITS) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + reference = tl.load(OUTPUT + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0) + contribution = tl.load( + CONTRIBUTION + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0 + ) + # count 0 (no captured reference) divides 0 by 0: NaN, never a silent value. + corrected = reference - contribution + tl.load(MASS + state) * (accumulator / normalizer) + tl.store(out, corrected, mask=dims < HEAD_DIM) + + +def _retention(device, key_retention): + """A float64 device scalar, built once; Triton would pass a Python float as fp32.""" + key = (device, key_retention) + if key not in RETENTION_SCALARS: + RETENTION_SCALARS[key] = torch.full((1,), key_retention, dtype=torch.float64, device=device) + return RETENTION_SCALARS[key] + + +def _build_capture_workspace(device, rows, heads, capacity, block_d): + empty = lambda *size, dtype=torch.float32: torch.empty(size, device=device, dtype=dtype) # noqa: E731 + return ( + empty(rows, heads, capacity), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + # Per-part radix histograms, the state before each pass (and the final + # threshold), per-part keys above the threshold and output offsets, the + # merged lse, the last-part tickets (zero between calls), and the + # contribution parts. + empty(rows, heads, SELECT_PARTS, 4, 256, dtype=torch.int32), + empty(rows, heads, 5, 2, dtype=torch.int64), + empty(rows, heads, SELECT_PARTS, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS, 2, dtype=torch.int32), + empty(rows, heads), + torch.zeros(rows, heads, device=device, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS), + empty(rows, heads, SELECT_PARTS, block_d), + ) + + +def _capture_workspace(device, rows, heads, capacity, block_d): + """Capture scratch reused across layers. + + Captures issued on one stream run in order, so each layer may overwrite + the previous layer's scratch; one entry per stream keeps memory bounded. + """ + stream = _current_stream(device) + shape = (rows, heads, capacity, block_d) + entry = CAPTURE_SCRATCH.get((device, stream)) + if entry is None or entry[0] != shape: + entry = (shape, *_build_capture_workspace(device, rows, heads, capacity, block_d)) + CAPTURE_SCRATCH[(device, stream)] = entry + return entry[1:] + + +def _consume_workspace(device, rows, heads, block_d): + """Consume split parts, grown per stream and reused across layers.""" + stream = _current_stream(device) + entry = CONSUME_SCRATCH.get((device, stream)) + if ( + entry is None + or entry[0].shape[0] < rows + or entry[0].shape[1] != heads + or entry[2].shape[-1] != block_d + ): + empty = lambda *size: torch.empty(size, device=device, dtype=torch.float32) # noqa: E731 + entry = ( + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + ) + CONSUME_SCRATCH[(device, stream)] = entry + return (None, *entry) + + +def supports(plan, q, dense_attention) -> bool: + """Whether the fused path covers this submission; otherwise use the batched path.""" + modes = set(plan.modes) + return ( + q.is_cuda + and q.shape[-1] <= 256 + and plan.key_retention < 1.0 + and (modes == {CONSUME} or (modes == {CAPTURE} and dense_attention is not None)) + ) + + +def capture_workspace(device, rows, heads, capacity, dim): + """Capture scratch for at most ``rows`` rows, owned by the caller. + + A CUDA graph bakes the scratch addresses in, so its owner keeps one + workspace alive for every bucket instead of the shape-keyed cache. + """ + return _build_capture_workspace(device, rows, heads, capacity, triton.next_power_of_2(dim)) + + +def fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention, + workspace=None, +): + """Capture or consume one uniform-mode submission without host readbacks. + + The kernels also write the store's ``valid``/``uses`` marks, so the host + only checks the plan's liveness instead of launching the bookkeeping. + ``workspace`` comes from ``capture_workspace`` and must cover the rows + and table width of the submission; consume uses only its split parts. + """ + store = plan.store + store._require_live(plan) + rows, heads, dim = q.shape + page = key_cache.shape[1] + groups = heads // key_cache.shape[2] + block_d = triton.next_power_of_2(dim) + scale = dim**-0.5 + for cache in (key_cache, value_cache): + if cache.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous head dimension") + if q.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous query head dimension") + selected = store.selected + if not selected.is_contiguous(): + raise ValueError("fused sparse attention derives index offsets from a contiguous pool") + if plan.modes[0] == CAPTURE: + # With the Triton backend, the dense split kernel also keeps the scores + # (and the first radix histograms), the dense merge runs inside the + # first selection pass, and selection runs in parallel parts. Other + # backends keep their call plus one serial selection kernel. + # Engines pass their backend adapter, which wraps the Triton kernel; + # callers may also pass the Triton function itself. + kernel = getattr(dense_attention, "kernel", dense_attention) + split = isinstance(kernel, TritonPagedAttention) or kernel is triton_paged_attention + if split: + output = torch.empty_like(q) + else: + output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + if rows == 0: + return output + # The round's page tables bound every row's context on the host. + capacity = block_tables.shape[1] * page + top = min(retained_count(capacity, plan.key_retention), store.key_bucket or capacity) + store.ensure_key_capacity(top) + selected = store.selected + if workspace is None: + workspace = _capture_workspace(q.device, rows, heads, capacity, block_d) + elif workspace[0].shape[1:] != (heads, capacity) or workspace[0].shape[0] < rows: + raise ValueError("capture workspace does not cover the submission") + ( + scores, + part_m, + part_l, + part_o, + seghist, + state, + greater, + offsets, + lse, + tickets, + part_mass, + part_c, + ) = workspace if workspace[0].shape[0] == rows else (buffer[:rows] for buffer in workspace) + if split: + folded = _PASS0_IN_SPLIT + retention = _retention(q.device, plan.key_retention) + _launch( + _capture_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + part_m, + part_l, + part_o, + scores, + seghist, + retention, + state, + tickets, + top, + capacity, + *q.stride(), + *key_cache.stride(), + *value_cache.stride(), + block_tables.stride(0), + HEAD_DIM=dim, + HEAD_GROUPS=groups, + NUM_HEADS=heads, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + HISTOGRAM=folded, + PARTS=SELECT_PARTS, + num_warps=4, + ) + first_pass = 1 if folded else 0 + for level in range(first_pass, 4): + _launch( + _select_pass_kernel, + (rows, heads, SELECT_PARTS), + part_m, + part_l, + part_o, + output, + scores, + context_lengths, + retention, + plan.slot_ids, + store.counts, + store.output, + lse, + seghist, + state, + tickets, + greater, + offsets, + layer, + top, + capacity, + *output.stride(), + PASS=level, + MERGE=level == first_pass, + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + SPLITS=NUM_SPLITS, + num_warps=_SCAN_WARPS, + ) + _launch( + _select_write_kernel, + (rows, heads, SELECT_PARTS), + scores, + context_lengths, + retention, + plan.slot_ids, + state, + offsets, + selected, + layer, + top, + capacity, + selected.shape[-1], + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + num_warps=_SCAN_WARPS, + ) + _launch( + _select_contribution_kernel, + (rows, heads, SELECT_PARTS), + scores, + value_cache, + block_tables, + context_lengths, + retention, + plan.slot_ids, + lse, + selected, + part_mass, + part_c, + tickets, + store.mass, + store.contribution, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + BLOCK_D=block_d, + BLOCK_T=_GATHER_BLOCK, + PARTS=SELECT_PARTS, + num_warps=4, + ) + return output + _launch( + _capture_kernel, + (rows, heads), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + _retention(q.device, plan.key_retention), + plan.slot_ids, + output, + scores, + selected, + store.counts, + store.mass, + store.output, + store.contribution, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=64, + num_warps=4, + ) + return output + output = torch.empty_like(q) + if rows: + if workspace is None: + workspace = _consume_workspace(q.device, rows, heads, block_d) + elif workspace[1].shape[0] < rows or workspace[1].shape[1] != heads: + raise ValueError("consume workspace does not cover the submission") + part_m, part_l, part_o = ( + workspace[1:4] if workspace[1].shape[0] == rows else (b[:rows] for b in workspace[1:4]) + ) + _launch( + _consume_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan.slot_ids, + selected, + store.counts, + part_m, + part_l, + part_o, + layer, + block_tables.stride(0), + selected.shape[-1], + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + num_warps=4, + ) + _launch( + _consume_merge_kernel, + (rows, heads), + part_m, + part_l, + part_o, + context_lengths, + plan.slot_ids, + store.mass, + store.output, + store.contribution, + output, + layer, + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + SPLITS=NUM_SPLITS, + num_warps=4, + ) + return output diff --git a/vllm_rlt/kernels/sparse_workspace.py b/vllm_rlt/kernels/sparse_workspace.py new file mode 100644 index 00000000..b365e352 --- /dev/null +++ b/vllm_rlt/kernels/sparse_workspace.py @@ -0,0 +1,37 @@ +"""Module-level scratch caches of the fused sparse decode kernels. + +The kernels reuse their scratch across layers of one stream, so the buffers live +outside any single call. They are keyed by ``(device, stream)`` so two engines on +different streams do not share them — but nothing owns an entry, which means a +closed engine's tensors stay reachable for the life of the process. The capture +score buffer alone is ``rows * heads * capacity`` floats, so a long-context run +leaves a large allocation behind. + +This module is deliberately free of Triton imports: the engine that owns the +stream has to be able to release the scratch on a CPU-only install too. +""" + +import torch + +# Scratch of one capture launch, per (device, stream). +CAPTURE_SCRATCH: dict = {} + +# Scratch of one consume launch, per (device, stream). +CONSUME_SCRATCH: dict = {} + +# ``key_retention`` as a float64 device scalar, per (device, retention). +RETENTION_SCALARS: dict[tuple[torch.device, float], torch.Tensor] = {} + + +def release_workspace(device=None) -> None: + """Drop the scratch cached for ``device``, or for every device. + + Call this only after the caller has drained the stream, or a kernel may + still be reading the tensors. Releasing too much is never a correctness + problem for the next user — an entry is rebuilt lazily on demand — only an + extra allocation, so tearing down a whole device is safe when several + engines share it. + """ + for cache in (CAPTURE_SCRATCH, CONSUME_SCRATCH, RETENTION_SCALARS): + for key in [key for key in cache if device is None or key[0] == device]: + del cache[key] diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 6b708df2..ceae6a14 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 @@ -82,7 +83,7 @@ def __init__(self, runner, scheduler): @dataclass -class SpeculativeRoundTicket: +class AsyncSpeculativeOutput: batch: SchedulerOutput buffer: SpeculativeRoundBuffer states: tuple[SpeculativeRequestState, ...] @@ -92,6 +93,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 +153,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 @@ -170,6 +178,8 @@ def __init__( copy_stream, rng: RngRegistry | None = None, runtime=None, + reference_store=None, + sparse_decode_config=None, ): super().__init__( model, @@ -180,6 +190,11 @@ 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_config = sparse_decode_config sizes = speculative_capture_sizes(execution, scheduler, config) self.graphs = ( DeviceRecurrentGraphs(model, cache, execution, False, capture_sizes=sizes) @@ -230,7 +245,11 @@ def _state(self, request, buffer): 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, ) if self.config.mode == "pipeline": width = min( @@ -264,7 +283,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 +301,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 +387,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_config.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 +434,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 +455,27 @@ 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) + 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 +487,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 +527,39 @@ 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, + ) + 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: @@ -489,7 +579,7 @@ def submit(self, batch): if self.config.mode != "pipeline": self.stats.drafted_tokens += sum(s.token_count - 1 for s in scheduled_requests) self.stats.verified_rows += count - return SpeculativeRoundTicket( + return AsyncSpeculativeOutput( batch, buffer, states, @@ -497,6 +587,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): @@ -563,7 +655,7 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer): p.tokens.copy_(torch.where(newborn, token, p.tokens)) p.depths.masked_fill_(newborn, 0) p.active.copy_((surviving | newborn) & (state.position + ids < limit)) - # Per-request wire order matches SpeculativeRoundTicket.collect(): + # Per-request wire order matches AsyncSpeculativeOutput.collect(): # prediction, accepted, emitted, next_position, drafted, verified. # A pipeline step verifies at most one token per request, so the # verified count equals emitted (0 or 1). @@ -685,6 +777,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..58accae1 --- /dev/null +++ b/vllm_rlt/worker/attention_reference.py @@ -0,0 +1,480 @@ +"""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 lease 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 owned by the lease that reserved it, and the lease object is the +identity: reserve binds the slot to a new lease, release requires that +same lease, and a plan is live only while every slot still points at the lease +it was built from. A slot handed to another request is bound to a different +lease, so the previous owner's plan cannot pass. This relies on two properties +the engine already guarantees: a request id is never reused while its previous +token is retiring, and slots are recycled only after their last reader +completed. +""" + +from dataclasses import dataclass + +import torch + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan +from vllm_rlt.kernels.sparse_workspace import release_workspace + +__all__ = ["ReferenceLease", "SparseDecodeReferenceStore", "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 + request_index: int + slots: tuple[int, ...] + slot_ids: torch.Tensor + row_index: torch.Tensor + + def __len__(self) -> int: + return len(self.slots) + + +class SparseDecodeReferenceStore: + """Device-resident reference state for the live tokens of one stage. + + The store is sized by the engine, which knows how many tokens can be live + at once. It never allocates per request: a speculative round holds one slot + per verification row, and several rounds of the same request may be + outstanding at the same time, so the capacity must cover every outstanding + round. + """ + + def __init__( + self, + *, + slots, + layers, + heads, + head_dim, + key_bucket, + device, + stream=None, + initial_key_capacity=None, + implementation="batched", + ): + if slots < 1 or layers < 1 or heads < 1 or head_dim < 1: + raise ValueError("reference store dimensions must be positive") + if key_bucket is not None and key_bucket < 1: + raise ValueError("key_bucket must be positive when set") + limit = key_bucket or 1 + initial = limit if initial_key_capacity is None else initial_key_capacity + if not 1 <= initial <= limit: + raise ValueError("initial key capacity must be within the reference capacity") + # Conservatively include all growth generations: an asynchronous + # stream may not have finished copying earlier buffers at the next + # growth. The engine budgets this before allocating its KV cache. + capacity = initial + self._peak_selected_capacity = capacity + while capacity < limit: + grown = min(limit, capacity * 2) + self._peak_selected_capacity += grown + capacity = grown + self.slots = slots + self.layers = layers + self.heads = heads + self.head_dim = head_dim + self.key_bucket = key_bucket + self.device = torch.device(device) + self.stream = stream + # Which CUDA attention path consumes this store's plans. + self.implementation = implementation + shape = (slots, layers, heads) + self.selected = torch.zeros((*shape, initial), dtype=torch.int32, device=device) + self.counts = torch.zeros(shape, dtype=torch.int32, device=device) + self.mass = torch.zeros(shape, dtype=torch.float32, device=device) + self.output = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + self.contribution = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + # Usability is decided by counts (zero means no captured reference) and + # by lease ownership, so no per-slot validity, use-count, position or + # acceptance tensor is kept: nothing reads them back. + self.free = list(range(slots)) + # Slot ownership is the lease object itself. Object identity is exact: + # a recycled slot is bound to a different lease, so a plan built for the + # previous owner cannot match, and no per-slot counter is needed. + self._bindings: dict[int, ReferenceLease] = {} + # Retained counts are host-known: the capture decides them before it + # submits. Keeping them here lets a reader size its gather without a + # device readback, which RFC #156 section 6.2 forbids. + self._counts_host: dict[tuple[int, int], int] = {} + + @property + def bytes(self) -> int: + """Payload bytes held by the store, excluding allocator overhead.""" + return sum( + tensor.numel() * tensor.element_size() + for tensor in ( + self.selected, + self.counts, + self.mass, + self.output, + self.contribution, + ) + ) + + @property + def peak_capacity_bytes(self) -> int: + """Conservative payload bound including all in-flight growth generations.""" + index_bytes = self.slots * self.layers * self.heads * self._peak_selected_capacity * 4 + return self.bytes - self.selected.nbytes + index_bytes + + def ensure_key_capacity(self, count: int) -> None: + """Grow shared index storage on the attention owner stream. + + Existing plans borrow the store rather than a fixed index pointer. + Prior reads/writes and the copy execute on the same owner stream; + record_stream protects the old allocation until that copy completes. + The logical key bucket and retained count never shrink with allocation. + """ + limit = self.key_bucket or 1 + if count > limit: + raise ValueError(f"reference capacity {limit} cannot hold {count} retained keys") + current = self.selected.shape[-1] + if count <= current: + return + if self.device.type == "cuda" and self.stream is not None: + if torch.cuda.current_stream(self.device) != self.stream: + raise RuntimeError("reference growth must run on the attention owner stream") + capacity = current + while capacity < count: + capacity = min(limit, capacity * 2) + old = self.selected + grown = torch.zeros((*old.shape[:-1], capacity), dtype=old.dtype, device=old.device) + grown[..., :current].copy_(old) + if old.device.type == "cuda": + old.record_stream(torch.cuda.current_stream(old.device)) + self.selected = grown + + @property + def outstanding(self) -> int: + """Number of slots currently bound to live tokens.""" + return len(self._bindings) + + def _index(self, values) -> torch.Tensor: + """Host-known slot or row indices on the device, without a blocking copy. + + A pageable host-to-device copy waits for the stream to drain, which + serializes the host with every in-flight round; a pinned source lets + the copy queue behind the work already submitted. + """ + host = torch.tensor(values, dtype=torch.int64, pin_memory=self.device.type == "cuda") + return host.to(self.device, non_blocking=True) + + def reserve(self, request_id: str, request_index: int, rows: int) -> ReferenceLease: + """Take one slot per verification row of a request's submission. + + The lease is the slot's owner identity. A slot returned to the pool and + handed to another request is bound to a different lease object, so plans + that still reference the old lease cannot pass a liveness check. + """ + 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") + # Peek first: a failed device allocation must not lose the slots. They + # are removed from the pool only once the lease is fully built. + slots = tuple(reversed(self.free[-rows:])) + slot_ids = self._index(slots) + row_index = torch.arange(rows, dtype=torch.int64, device=self.device) + del self.free[-rows:] + lease = ReferenceLease( + request_id=request_id, + request_index=request_index, + slots=slots, + slot_ids=slot_ids, + row_index=row_index, + ) + for slot in slots: + self._bindings[slot] = lease + return lease + + 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 owner. + """ + if isinstance(leases, ReferenceLease): + leases = (leases,) + cleared = [] + for lease in leases: + for slot in lease.slots: + if self._bindings.get(slot) is not lease: + raise RuntimeError("attention reference slot released by a different owner") + del self._bindings[slot] + cleared.append(slot) + if cleared: + self.free.extend(cleared) + self._forget_counts(cleared) + if self.stream is None: + self._clear(self._index(cleared)) + return + # A ticket retires from the host's current stream. An index built + # there would be freed, and reusable by that stream, before the + # owner stream's index_fill_ reads it; build it on the owner stream. + with torch.cuda.stream(self.stream): + self._clear(self._index(cleared)) + + def discard(self) -> None: + """Free every binding without device work; the caller has drained.""" + self._bindings.clear() + self._counts_host.clear() + self.free = list(range(self.slots)) + + def close(self) -> None: + """Release the pool and the kernel scratch this store's device owns. + + The fused kernels cache per-stream scratch at module level, which + nothing else can reach. This store is the execution owner of the stream + those launches ran on, so it drops them here. The caller must have + drained the stream: the tensors must not be freed under a running + kernel. A later engine rebuilds them lazily. + """ + self.discard() + # The fused kernels cache per-stream scratch at module level, and this + # store is the execution owner of the stream they ran on. + release_workspace(self.device) + + 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 their owning leases 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, row_owners = [], [] + for owner, mode in zip(owners, modes): + if owner is None: + if mode != DENSE: + raise ValueError("a sparse row needs a reference owner") + slots.append(-1) + row_owners.append(None) + continue + lease, row = owner + slot = lease.slots[row] + if self._bindings.get(slot) is not lease: + raise RuntimeError("attention reference slot released before its plan") + slots.append(slot) + row_owners.append(lease) + return SparseDecodePlan( + modes=tuple(modes), + slots=tuple(slots), + owners=tuple(row_owners), + key_retention=key_retention, + slot_ids=self._index(slots), + 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, owner, mode in zip(plan.slots, plan.owners, plan.modes): + if slot == -1 and mode == DENSE: + continue + if self._bindings.get(slot) is not owner: + raise RuntimeError("stale attention reference plan for a retired token") + + def note_count(self, slot: int, layer: int, count: int) -> None: + """Record the retained count of one capture row on the host. + + The count is decided by the capture, so it never needs a device + readback; ``host_count`` is the reader-side counterpart. + """ + self._counts_host[(slot, layer)] = count + + def _forget_counts(self, slots) -> None: + """Drop the host counts of host-known slots, so none sizes a later read.""" + for slot in slots: + for layer in range(self.layers): + self._counts_host.pop((int(slot), layer), None) + + def host_count(self, slot: int, layer: int) -> int | None: + """Retained keys of ``(slot, layer)``, or ``None`` before a capture.""" + return self._counts_host.get((slot, layer)) + + def _row_slots(self, plan: SparseDecodePlan, rows: list[int]) -> torch.Tensor: + """Slot ids of ``rows`` without a host-to-device copy of a Python list. + + ``plan.slot_ids[rows]`` costs ~2.4 ms per call because the list becomes + a device tensor first; grouped rows are contiguous, so derive the index + on the device instead. + """ + if len(rows) == len(plan.slots): + return plan.slot_ids + if rows[0] + len(rows) == rows[-1] + 1: + index = torch.arange(rows[0], rows[0] + len(rows), device=self.device) + return plan.slot_ids.index_select(0, index) + return self._index([plan.slots[row] for row in rows]) + + def capture(self, plan: SparseDecodePlan, layer: int) -> None: + """Assert that a finished capture row still owns its slots. + + The attention kernel writes the retained state of a capture row and then + calls this on the same stream. Nothing is recorded: a reader decides + availability from the retained count, and ownership from the lease. + """ + self._require_live(plan) + + def consume(self, plan: SparseDecodePlan, layer: int) -> None: + """Assert that a finished consume row still owns its slots.""" + self._require_live(plan) + + def write_reference( + self, + plan: SparseDecodePlan, + layer: int, + row: int, + *, + selected, + counts, + mass, + output, + contribution, + ) -> None: + """Publish the retained state of one capture row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = selected.shape[-1] + self.ensure_key_capacity(count) + self.selected[slot, layer, :, :count] = selected.to(torch.int32) + self.counts[slot, layer] = counts.to(torch.int32) + self.note_count(slot, layer, int(count)) + self.mass[slot, layer] = mass + self.output[slot, layer] = output + self.contribution[slot, layer] = contribution + + def read_reference(self, plan: SparseDecodePlan, layer: int, row: int): + """Borrow the retained state of one consume row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = self.host_count(slot, layer) + if count is None: + # A fused capture decides its count on the device. + count = int(self.counts[slot, layer].max()) + if not count: + raise RuntimeError("a consume row has no captured reference") + return ( + self.selected[slot, layer, :, :count].long(), + self.mass[slot, layer], + self.output[slot, layer], + self.contribution[slot, layer], + ) + + def accept(self, leases, accepted: torch.Tensor) -> None: + """Invalidate the references of rejected rows on the device. + + ``accepted`` is the number of accepted candidates of every request of + the submission. A verification row of a request survives exactly when + its row index is at most the accepted count: the root row is always + kept, and a rejected candidate invalidates its dependent descendants. + The comparison stays on the device so that acceptance never requires a + host readback. + + Retirement is memory-bounded: it only clears the per-row metadata that + gates readers (``counts`` and ``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) + # A boolean-mask gather of the rejected slots would read its size back + # to the host. Masked writes over the round's own rows stay on device; + # they touch only these rows' per-layer metadata, never the index bucket. + keep = rows <= limit + self.counts[slot_ids] = torch.where(keep[:, None, None], self.counts[slot_ids], 0) + self.mass[slot_ids] = torch.where(keep[:, None, None], self.mass[slot_ids], 0) + # Which rows were rejected is known only on the device, so drop the host + # counts of every row of these leases: a later reader then sizes from + # the device counts, which are zero for rejected rows, and a stale host + # count can never size a gather. + self._forget_counts(slot for lease in leases for slot in lease.slots) + + def _clear(self, slots: torch.Tensor) -> None: + """Retire the metadata of the given slots, ordered on the owner stream. + + Only the tensors that gate readers are cleared. The retained index, + output and contribution buffers stay allocated and are overwritten by + the next capture's used prefix, so retiring a slot never writes the + whole bucket (which would be tens of GiB per batch at full retention). + """ + + def zero(): + # index_fill_ passes the scalar to the kernel; assigning a Python + # scalar through an index tensor copies it from the host and blocks. + for tensor in ( + self.mass, + self.counts, + ): + tensor.index_fill_(0, slots, 0) + + if self.stream is not None: + with torch.cuda.stream(self.stream): + zero() + else: + zero() + + def invalidate(self, slots) -> None: + """Invalidate the references of the given slots without freeing them.""" + if len(slots) == 0: + return + if isinstance(slots, torch.Tensor): + # Callers hand host lists; a device tensor costs one readback here. + self._forget_counts(slots.tolist()) + slots = slots.to(device=self.device, dtype=torch.int64) + else: + self._forget_counts(slots) + slots = self._index(list(slots)) + for tensor in ( + self.mass, + self.output, + self.contribution, + self.selected, + self.counts, + ): + tensor.index_fill_(0, slots, 0) diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index 8f2c020b..03880166 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -39,6 +39,7 @@ class KVMetadata: writable: bool cu_seqlens_q: torch.Tensor | None = None max_seqlen_q: int = 1 + reference_plan: object | None = None @dataclass(frozen=True, eq=False) @@ -62,6 +63,10 @@ class DeviceKVMetadata: context_lengths: torch.Tensor # Inactive pipeline slots must not overwrite any live KV address. write_mask: torch.Tensor | None = None + # Sparse decode attention plan of this submission: one row mode and + # reference slot per row. Host-known; the attention layer borrows it and + # the runner-owned store keeps the referenced state alive. + reference_plan: object | None = None # The device-owned path currently uses independent single-query rows. # These layout fields also let recurrent graphs stage the same contract. cu_seqlens_q: torch.Tensor | None = None @@ -100,6 +105,7 @@ def prepare_device( block_tables, context_lengths, write_mask=None, + reference_plan=None, ) -> DeviceKVMetadata: """Borrow stream-ordered metadata without reading device positions on CPU. @@ -115,7 +121,10 @@ def prepare_device( block_tables=block_tables, context_lengths=context_lengths, write_mask=write_mask, + reference_plan=reference_plan, ) + if reference_plan is not None and len(reference_plan.modes) != len(position_ids): + raise ValueError("reference plan needs one row mode per query row") if write_mask is not None and ( write_mask.shape != position_ids.shape or write_mask.dtype != torch.bool diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..d9d92cca 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -7,6 +7,7 @@ import torch +from vllm_rlt.attention.sparse import CAPTURE, DENSE, reference_mode from vllm_rlt.config import ExecutionConfig, ExitConfig, SchedulerConfig from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import ExecutionBatch, Operation, Request @@ -160,9 +161,14 @@ def __init__( scheduler_config=None, staged_prefill=None, rng: RngRegistry | None = None, + reference_store=None, + sparse_decode_config=None, ): self.model = model.eval() self.cache_manager = cache_manager + self.reference_store = reference_store + self.sparse_decode_config = sparse_decode_config + self.attention_references = {} self.metadata = cache_manager.metadata_builder self.runtime = ModelRuntime(self.model, cache_manager) parameter = next(model.parameters()) @@ -1016,6 +1022,8 @@ def _prepare_forward(self, prepared): return replace(prepared, forwards=(PreparedForward(self.runtime.prepare(forward)),)) hidden = self._load_hidden(batch, prepared, workspace, size) if batch.operation == Operation.CODA: + for request in requests: + self._release_attention_reference(request.request_id) forward = ForwardBatch( Operation.CODA, hidden=hidden[: batch.num_rows] if kv_metadata else hidden, @@ -1027,6 +1035,8 @@ def _prepare_forward(self, prepared): kv = prepared.kv if kv is not None: metadata = (kv,) + metadata + if self.reference_store is not None: + metadata = tuple(self._prepare_sparse_reference(batch, kv) for kv in metadata) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden[: batch.num_rows] @@ -1040,6 +1050,53 @@ def _prepare_forward(self, prepared): prepared, kv=metadata[0], forwards=(PreparedForward(self.runtime.prepare(forward)),) ) + def _release_attention_reference(self, request_id): + state = self.attention_references.pop(request_id, None) + if state is not None: + # release() creates the index tensor used by the metadata clears. + # Allocate it on core too, so its allocator lifetime follows the + # clear kernels even when CODA or teardown calls from another stream. + with ( + torch.cuda.stream(self.core_stream) + if self.core_stream is not None + else nullcontext() + ): + self.reference_store.release(state[2]) + + def _prepare_sparse_reference(self, batch, kv): + """Borrow token-owned references; prompt and padded rows remain dense. + + Preparation runs on core, after the request's preceding event. CODA + retires the reference after its final recurrent read; request teardown + waits for that same event before releasing an unfinished token. + """ + requests = [span.request for span in batch.spans for _ in range(span.token_count)] + owners, modes = [], [] + for request, (_, depth, position) in zip(requests, kv.rows): + mode = reference_mode(depth) if position >= len(request.prompt_token_ids) else DENSE + state = self.attention_references.get(request.request_id) + if state is not None and (state[0] is not request or state[1] != position): + self._release_attention_reference(request.request_id) + state = None + if mode == CAPTURE: + if state is None: + lease = self.reference_store.reserve(request.request_id, 0, 1) + state = (request, position, lease) + self.attention_references[request.request_id] = state + if state is None: + mode = DENSE # No reference loop for this token: stay dense. + owners.append(None if mode == DENSE else (state[2], 0)) + modes.append(mode) + if all(mode == DENSE for mode in modes): + return kv + padding = len(kv.position_ids) - len(owners) + owners.extend([None] * padding) + modes.extend([DENSE] * padding) + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode_config.key_retention + ) + return replace(kv, reference_plan=plan) + def _load_hidden(self, batch, prepared, workspace, size): if prepared is not None and prepared.batch_inputs is not None: return prepared.batch_inputs.hidden @@ -1269,6 +1326,8 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu ) except Exception: self._synchronize_submission(stream, prepared) + for span in batch.spans: + self._release_attention_reference(span.request.request_id) # A failed preparation may advance a pool without enqueueing a batch. # Retire all prior readers before any pool can wrap after that gap. self.submissions.drain() @@ -1364,6 +1423,7 @@ def _free_runner_slot(self, request_id): event = self.events.pop(request_id, None) if event is not None: event.synchronize() + self._release_attention_reference(request_id) self.prompt_token_owners.pop(request_id, None) if self.async_state is not None: self.async_state.release(request_id) diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 85f0e235..21e2008d 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -9,6 +9,7 @@ import torch +from vllm_rlt.attention.sparse import DENSE, reference_mode from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import Operation, Request from vllm_rlt.worker.buffers import stage_metadata @@ -152,6 +153,8 @@ def __init__( *, runtime=None, compute_stream=None, + reference_store=None, + sparse_decode_config=None, ): self.model, self.cache, self.config = model, cache, config self.metadata = cache.metadata_builder @@ -159,6 +162,8 @@ def __init__( self.pipeline_states: dict[str, PipelineState] = {} self.device = next(model.parameters()).device self.compute_stream = compute_stream + self.reference_store = reference_store + self.sparse_decode_config = sparse_decode_config self.submissions = SubmissionQueue() self.input_buffers = { dtype: StagedInputTensor(dtype, self.device, use_uva=execution.use_uva) @@ -175,9 +180,14 @@ def __init__( def _packed(self, packed): capabilities = self.cache.attention_capabilities - return packed and capabilities.packed_prefill and capabilities.packed_matches_decode + return ( + packed + and self.reference_store is None + and capabilities.packed_prefill + and capabilities.packed_matches_decode + ) - def _prepare_inputs(self, batch): + def _prepare_inputs(self, batch, references=()): """Plan all host-known inputs, then publish one snapshot per dtype.""" for buffer in self.input_buffers.values(): buffer.reset() @@ -219,8 +229,33 @@ def metadata(ids, positions, depths, *, packed=False): for active, layouts in steps ] verification = tuple(self.metadata.bind(kv, uploaded) for kv in verification) + if references: + steps = [ + ( + active, + self._prepare_sparse_reference( + layouts, [(references[i], offset) for i in active] + ), + ) + for offset, (active, layouts) in enumerate(steps) + ] + owners = [(lease, row) for lease in references for row in range(len(lease))] + verification = self._prepare_sparse_reference(verification, owners) return tokens, steps, verification + def _prepare_sparse_reference(self, layouts, owners): + result = [] + for kv in layouts: + modes = [reference_mode(depth) for _, depth, _ in kv.rows] + if all(mode == DENSE for mode in modes): + result.append(kv) + continue + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode_config.key_retention + ) + result.append(replace(kv, reference_plan=plan)) + return tuple(result) + def _recurrent(self, hidden, metadata): forward = ForwardBatch( Operation.RECURRENT, hidden=hidden, metadata=metadata, graphs=self.graphs @@ -257,11 +292,27 @@ def submit(self, batch): if self.compute_stream is not None else nullcontext() ): - if self.config.mode == "pipeline": - submission = self._submit_pipeline(batch) - else: - inputs, steps, verification = self._prepare_inputs(batch) - submission = self._submit_round(batch, inputs, steps, verification) + references = [] + try: + if self.reference_store is not None: + for index, span in enumerate(batch.requests): + references.append( + self.reference_store.reserve( + span.request.request_id, index, span.token_count + ) + ) + if self.config.mode == "pipeline": + submission = self._submit_pipeline(batch) + else: + inputs, steps, verification = self._prepare_inputs(batch, references) + submission = self._submit_round(batch, inputs, steps, verification) + finally: + if references: + # Every candidate has finished its final recurrent read. + # No reference survives this block round, including the + # rejected suffix. Clear/reuse is ordered on core after + # those readers, without waiting for host acceptance. + self.reference_store.release(references) self.submissions.complete(submission.event) return submission except Exception: