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