[0z5a]

Summary

Adds WISE working-set decode (Wan and Chen, arXiv 2609.27373) for recurrent-depth models, opt-in through wise_config. WISE finds 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 at depths 2–3, with the paper's B=32 blocks and η=0.95 selection rule.

This is a cached-decode variant. Prefill stays dense, and earlier tokens' K/V come from decode under their own supports. The paper's reference recomputes the whole prefix instead.

On one GH200 (Ouro-1.4B, bf16, async + CUDA graphs), the best configuration is FlashAttention-4 prefill + the Triton WISE decode kernel. Its decode throughput is 1.08–1.19× (1K), 1.16–1.39× (4K) and 1.23–1.41× (8K) that of the best dense configuration, and answer quality on 200 HotpotQA/2Wiki prompts is unchanged within noise (details below).

What changes

• vllm_rlt/kernels/wise_attention.py: split-KV Triton paged decode with per-row modes:

  • dense;
  • discover: dense, while recording each 32-token tile's max and sum;
  • sparse: only 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. Programs cover about eight tiles with one warp each; one long dense row otherwise sets the latency of a launch whose other rows are sparse.

• vllm_rlt/kernels/wise_flash.py: FlashAttention 2/3/4 decode.

  • Sparse rows attend the per-layer head union of their support, at page granularity, through a compacted page table.
  • Discovery rereads K once to rebuild block mass from FlashAttention's LSE.
  • WiseConfig(kernel="triton") keeps FlashAttention for prefill and dense decode and decodes WISE with the Triton kernel. This is the fastest configuration measured.

• vllm_rlt/wise.py: WiseConfig and WiseState. WiseState holds per-request slots and a [slot, layer, head, block] support bitmask. Support is tagged with its (request, position): reuse never uses support from another position or an incomplete discovery window, and slots are freed with the allocation. A token whose discovery is incomplete (a discovery depth that ran in a packed mixed batch, or a KV offload between discovery and reuse) attends densely at its remaining depths; WiseState.fallbacks counts those rows.

• Integration:

  • The model runner marks the decode rows of decode traversals. Prefill-chunk rows mixed into a batch, and padding rows, stay dense.
  • CUDA graphs key WISE and dense traversals separately and own their WISE inputs.
  • Async scheduling works unchanged.
  • StageEngine / LLM take wise_config, and 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.

• benchmarks/wise_{throughput,kernel_bench,quality}.py.

With wise_config=None (the default) nothing changes.

Tests

tests/test_wise.py (GPU; the module skips without Triton, so CPU CI only collects it):

• Kernel: the modes agree with each other, and one split at four warps is bitwise the existing dense kernel. The fused selection matches the reference rule and the union across discovery depths. Sparse output matches masked attention.
• FlashAttention: compaction with every block selected is bitwise dense FlashAttention. LSE-derived block mass matches softmax block sums, and union-sparse output matches its mask (FA2/FA3/FA4).
• Engine: all-blocks WISE generates the same tokens as dense decode, eager and with CUDA graphs. Async WISE matches synchronous WISE. The Triton WISE kernel works beside FlashAttention prefill.
• Speculation: a WISE draft leaves greedy output unchanged (Triton, FA3, FA4, eager and graphs).
• Safety: reuse without this token's complete discovery falls back to dense attention, covered for offload/resume and for FA4 mixed packed batches with and without graphs.

Full suite with --run-gpu on one GH200 (FA2/FA3/FA4 installed):

• This branch: 1116 passed, 8 failed.
• main (ff05fe1): 1094 passed, 9 failed.

The failures on this branch are a subset of main's. 6 are PD tests that need a second GPU, and 2 are FlashAttention tests that fail in this environment on main too. main's ninth failure, a CUDA memory-plan test, passes here. The 22 extra passes are the 21 WISE tests plus that test. Ruff 0.13.2 lint and format are clean.

Results (one GH200, Ouro-1.4B bf16)

Measured at 190f36c with torch 2.14.1+cu130, Triton 3.8.0, FA3 3.0.1 and FA4 4.0.0b30. Setup: C concurrent requests of exactly L real prompt tokens (Ouro-tokenized HotpotQA/2Wiki text), greedy, 256 output tokens with ignore_eos, always 4 loops. Decode throughput is timed inside the engine step loop after an untimed warm-up, and counts only tokens generated after every request has its first token. Every arm runs in a fresh process with an identical KV pool, async scheduling and CUDA graphs. Two repeats each.

Decode tokens/s:

L × C	FA3 dense	FA4 dense	WISE, Triton	WISE, FA4 prefill + Triton decode	WISE / best dense
1K × 64	1350 / 1406	1382 / 1375	1568 / 1514	1538 / 1608	1.08–1.19×
4K × 20	424 / 426	411 / 410	521 / 498	493 / 571	1.16–1.39×
8K × 10	215 / 211	201 / 197	277 / 272	265 / 268	1.23–1.41×

The ratio range spans both repeats of every arm. The FA4-prefill WISE arm's two repeats differ by up to 16% (4K), so the upper ends are noisy. An earlier version of this table, measured before the review fixes in 853f4f2, counted some tokens generated before the decode timer started (0.4–0.9% here) and reported the density of individual discovery selections (0.22 / 0.13 / 0.10) instead of the reused union.

Mean reuse-depth support density (the union the reuse depths attend) is 0.28 at 1K, 0.18 at 4K and 0.15 at 8K. WISE accelerates decode only, so end-to-end throughput including prefill (FA4 prefill in both arms) moves less: 4K 231–232 → 225–273 tokens/s, 8K 116–120 → 117–127.

Per-call kernel time from CUDA-graph replay (1K × 64, fixed support density 0.13 set by the kernel benchmark):

Kernel	Dense	Sparse
Triton	234 µs	51 µs
FA3	169 µs	59 µs
FA4	170 µs	58 µs

At 4K × 20, a mixed batch of half discovery rows and half sparse rows takes 143 µs, against 211 µs for dense FA3.

Answer quality through the engine: 200 fixed HotpotQA/2Wiki plain-completion prompts, greedy, 32 tokens, first line scored with the WISE authors' benchmark_scores. Both arms use FA4 prefill, async and graphs.

Subset	Dense F1	WISE F1	ΔF1 [95% bootstrap CI]	EM (dense / WISE)
HotpotQA (100)	0.4617	0.4545	−0.0072 [−0.0179, +0.0010]	34 / 32
2Wiki (100)	0.3342	0.3367	+0.0025 [−0.0064, +0.0116]	24 / 26
All (200)	0.3980	0.3956	−0.0024 [−0.0093, +0.0039]	58 / 58

11 of 200 first-line answers change.

Limits

• The FlashAttention-only path (discovery also on FlashAttention) is slower than dense FlashAttention. FlashAttention cannot report block mass, so discovery costs an extra K read; use kernel="triton".
• FA2 cannot be captured in CUDA graphs here, and its ≥256-token pages make the head-union support nearly dense.
• A WISE draft is lossless with high acceptance (0.98–0.995) but slower than dense decode: drafting runs every loop.
• Quality was checked on short answers (32 tokens). Long generations are not covered.
• Not supported: async speculative decoding with a WISE draft, and Huginn.

[hsliuustc0106]

Request changes. Opt-in support reuse addresses repeated attention scans, but two supported execution paths lose required discovery state. Dense decode is a credible fallback when discovery cannot be preserved. The inline findings also identify two benchmark-accounting issues affecting the reported decode throughput and reuse density.

Reviewed the complete 15-file diff, relevant callers and state ownership, graph/async integration, speculative verification, tests, benchmark accounting, and the authors' support-selection rule. Reviewed head: 5c36237adeac85fbede5660058249be3ed703bc4; base and merge base: ff05fe15d182af97a30be321c99e20d5d3b17b92. The head was rechecked unchanged before submission.

Validation: existing CI passed lint/format and 803 tests, with 379 skipped. The CI merge commit has the same tree as the reviewed head. No PR code or GPU work was executed for this review; skipped device tests do not validate the WISE kernels. Contributor-reported GH200 results remain independently unverified, and quality evidence covers only 200 short-answer prompts.

[hsliuustc0106]

[P1] Preserve WISE state across KV offload

With preemption enabled, capacity offload after the discovery depths calls free(), which now deletes the request's support slot and discovery tags. PreemptionManager.resume() restores the KV snapshot and hidden state while retaining the saved recurrent progress, but it does not restore WISE state. A request resumed at depth 2 therefore raises WISE reuse without this position's complete discovery instead of continuing. Snapshot and restore the support and tags alongside the recurrent state, or explicitly reject this unsupported combination, and cover offload/resume between discovery and reuse.

[hsliuustc0106]

[P1] Handle discovery in mixed packed FA4 batches

With FlashAttention-4, SchedulerConfig(mode="mixed"), and ExecutionConfig(atomic_prefill=False), decode rows share packed metadata with prompt-chunk rows. _prepare_forward() marks the decode rows, but the cu_seqlens_q branch above returns dense prefill attention before WISE records their discovery. Graph selection likewise treats the packed batch as dense. If a discovery depth runs in such a mixed batch and a later depth runs in an ordinary decode batch, WiseState.plan() raises its discovery-order or incomplete-discovery error. Process discovery for the marked rows, or retain a dense fallback for that token's remaining depths; cover the mixed-to-decode transition with and without graphs.

[hsliuustc0106]

[P2] Count only decode work inside the timed interval

Refill scheduling lets early requests decode while later prompts are still prefilling. timed() starts its decode timer once every request has at least one generated token, but concurrency * (output_tokens - 1) assumes every request has generated exactly one token at that boundary. For cohorts requiring multiple prefill batches, tokens already generated by the early requests are counted in the numerator without their execution time. This inflates decode throughput and can distort comparisons, particularly when prefill backends differ. Count actual work after the timing boundary, or finish the whole cohort's prefill before timing decode.

[hsliuustc0106]

[P2] Distinguish discovery density from reused union density

For MODE_DISCOVER_UNION, chosen is ORed with the previous support, but this counter still adds the pre-union keep_count. Consequently, mean_density(), which the throughput benchmark emits as support_density, averages individual discovery selections rather than measuring the reuse-depth union density reported in the PR. For example, two disjoint one-block selections over two causal blocks report 50%, although reuse attends 100%. Measure the final union after discovery completes, or relabel the metric and the associated claims.

[0z5a]

@hsliuustc0106 Thank you for the review. I pushed 853f4f2. It takes the dense fallback you suggested for both P1 cases:

1. KV offload between discovery and reuse. free() still releases the support slot and tags. When the resumed request reaches a reuse depth without its token's complete discovery, that row now attends densely instead of raising. The next token discovers again. New test: test_offload_between_discovery_and_reuse_resumes_densely stops a request after the discovery depths, offloads its KV through scheduler.preempt, resumes it, and checks that the output matches the dense engine.

2. Mixed packed FA4 batches. A discovery depth that runs in a packed batch records no support. When a later depth of that token runs in an ordinary decode batch, it now attends densely for the rest of that token, both for a later discovery depth and for the reuse depths. New test: test_mixed_packed_fa4_decode_rows_fall_back_to_dense (FA4, mode="mixed", atomic_prefill=False, with and without CUDA graphs). It checks that the all-blocks control matches the dense engine and that fallbacks actually happened.

   WiseState.fallbacks counts the rows that fell back. A unit test covers each transition: broken window, release, and stale position.

3. Decode throughput accounting. timed() now counts only the tokens generated after the decode timer starts (decode_tokens in the record). It also resets the density counters after the warm-up, which they previously included.

4. Support density. The counter now measures the union that the reuse depths attend. The first discovery depth adds its selection and the causal blocks. A union depth adds only the blocks it newly sets. In your example, two disjoint one-block selections over two blocks now report 100%. The kernel test checks the union's count.

Validation on one GH200 (torch 2.14.1, Triton 3.8.0, FA3 3.0.1, FA4 4.0.0b30; FA2 not installed):

• Core GPU suite at 853f4f2 (attention, CUDA graphs, KV cache, mixed execution, speculative, async speculative, split prefill, FlashAttention, backend selection, WISE): 338 passed, 5 skipped. The skips are by design: two FA2 WISE cases, two FA4-only packed-prefill cases, and one FA4 fp32-operand case. The same run with main's pre-capture gc.collect() from [Bugfix] Collect garbage before CUDA graph capture #141 added gave the same result.
  • An earlier run of this suite, before FA3 was installed, aborted inside a capture while garbage was being collected. That is the intermittent issue [Bugfix] Collect garbage before CUDA graph capture #141 fixes on main, and this branch's base ff05fe1 predates the fix.
• The five new or changed tests, run on the old head 5c36237 with only the test file replaced: 5 failed, with the errors you described. Three raise WISE reuse without this position's complete discovery, one raises WISE discovery out of order, and the union density check counts [229, 244] instead of [1, 0].

The throughput ratios, the support densities (0.22 / 0.13 / 0.10, and the 0.13 in the kernel table) and the quality-run density in the PR description were measured with the old accounting, so the density numbers overstate sparsity at reuse depths and the decode tok/s may be inflated. I have not rerun them yet and will update the description when I do. The kernel times themselves come from a separate benchmark and do not depend on these two counters.

[0z5a]

@hsliuustc0106 Follow-up to the two benchmark findings. I re-measured the throughput table with the corrected accounting at 190f36c. That commit only updates the benchmark docstring to describe the new decode count. The PR description now has the new numbers.

• Decode token count: tokens generated before the decode timer started were 0.4–0.9% of the decode tokens in these cases (for example 16,180 counted at 1K × 64 against 16,320 assumed before), so this fix moves decode tok/s by under 1%.
• Support density: the reused-union density is 0.28 / 0.18 / 0.15 at 1K / 4K / 8K. The previously reported values were 0.22 / 0.13 / 0.10, so the old numbers understated the work at reuse depths by roughly a third.
• Decode throughput, WISE over the best dense FlashAttention arm: 1.08–1.19× at 1K × 64, 1.16–1.39× at 4K × 20 and 1.23–1.41× at 8K × 10, on one GH200 with torch 2.14.1, Triton 3.8.0, FA3 3.0.1 and FA4 4.0.0b30. These come from the same four arms, prompts, cases and two repeats as before. The FA4-prefill WISE arm's repeats differ by up to 16%, so the upper ends are noisy.

The answer-quality results are unchanged: that path has no offload or mixed batches, so 853f4f2 does not change its outputs, and its table reports no density. The kernel table uses a fixed density set by the kernel benchmark.