[sysuTobo]

Summary

Adds an explicit flashinfer_fa2 attention backend for NVIDIA Ada (SM89: L20/L40S/RTX 40) and wires it through the engine, CLI/server and the GSM8K harness.

• FlashInfer FA2 paged decode: BF16, page size 16, head_dim 128, NHD, LAST_EXITED and SHARED.
• Explicit opt-in only (--attention-backend flashinfer_fa2); auto keeps FA4/FA3/FA2/Triton/Torch.
• Eager, CUDA Graphs and asynchronous scheduling. No packed prefill, PD or speculative paths; async speculative is rejected at selection (device_context_lengths stays undeclared).

Claimed as the Ouro x Ada cell in #69. #112 adds XQA for SM120, and XQA rejects SM89, so Ada needs the FA2 path.

Design notes

• KVCacheManager.prepare_attention / ModelRuntime.forward build the CSR metadata (indptr, compact indices, last_page_len) once per traversal. Construction is fixed-shape and device-only.
• The graph paths prime the graph entry's private metadata after _copy_graph_inputs; the first call of a traversal copies the staging CSR into the plan buffers (later layers skip it) and a capture records one copy group, so a replay consumes what prepare_step wrote without running Python inside the graph.
• The adapter uses FlashInfer's tensor-core CUDAGraphBatchDecodeWithPagedKVCacheWrapper. The CUDA-core wrapper bakes its work schedule into plan-time state: on SM89 a full-capacity plan replayed with a short live context stalls once the table reaches 64+ pages (reproduced for GQA and MHA).
• Split-KV is disabled. With splitting, a row's reduction order depends on the planned width and batch composition (verified: the same 1,500-token row differs between a 128- and 256-page plan, and alone vs inside a batch of eight). disable_split_kv makes eager, graph replay, plan width and batch composition bit-identical; the cost is low-batch long-context throughput (table below).
• Plans are keyed by (rows, q_heads, kv_heads, width bucket), with the width rounded up to a power of two. A 1K-token eager decode creates 6 plans instead of 26. workspace_size()/plan() receive pinned host inputs, so plan creation does not copy device metadata back to the host; plan_creations reports the count.
• Captured plans are pinned and eager plans are evicted first (plan cache bound 32). After the first plan the adapter requires resolved_backend == "fa2" and records it in info, so a patch release routing to cuDNN fails loudly.
• benchmarks.gsm8k run --cuda-graphs is recorded in the summary and does not change the frozen baseline fingerprint.

Validation (1xL20, torch 2.13.0, CUDA 13.0, FlashInfer 0.7.2.post1)

• CPU suite: 896 passed / 7 skipped (tests/test_serving.py cannot be collected on this image because its aiohttp predates web.AppKey).
• tests/test_flashinfer_fa2_attention.py: 22 cases (16 GPU, 6 CPU). Selection/rejection, strided GQA/MHA against Torch and Triton, padding tri-state, zero rows, batch composition, vectorized CSR, no host synchronization across a width-bucket change, bounded plan creation, bitwise width/batch/graph invariance at a 128-page plan, and tiny-Ouro eager/graphs/async parity including a 2048-token config.
• Real Ouro-1.4B greedy: Triton eager, FA2 eager, FA2 graphs and FA2 async+graphs produce identical tokens; plan_misses == 0; one capture, 60 replays.
• GSM8K-87 default protocol (HF baseline 59/87):

run	correct	accuracy	gate
Triton eager	61/87	70.11%	pass
FA2 eager	59/87	67.82%	pass (+0.00 pp)
FA2 + CUDA Graphs	59/87	67.82%	pass (+0.00 pp)

Eager and graphs are bit-identical (87/87 sequences), no runaway, mean 106.7 tokens (eager 717.8 s, graphs 279.6 s). Against Triton, 67/87 sequences match; four answers flip. The no-split arithmetic reproduces the earlier split-KV graph run on all 87 questions; the earlier split-KV eager run scored 60/87.

Performance (diagnostic, not a benchmark)

Matched fixed work: real Ouro-1.4B BF16, greedy depth 4, 64 output tokens, 32-token prompts, one wave of N requests, second run after a warmup that captures the graph. All five columns come from one session.

concurrency	Triton eager	Triton graphs	FA2 eager	FA2 graphs	FA2 async+graphs
1	5.20 s (12.3 tok/s)	1.52 s (42.2)	5.49 s (11.7)	1.59 s (40.2)	1.58 s (40.6)
8	5.34 s (95.8)	1.76 s (290.2)	5.55 s (92.2)	1.84 s (278.1)	1.81 s (282.8)
32	6.26 s (327.3)	2.20 s (930.1)	6.32 s (324.1)	2.24 s (915.4)	2.15 s (952.3)

Eager FA2 is at parity with eager Triton. CUDA Graphs remove most per-step launch overhead for both backends (2.8-3.5x over eager on this small model); Triton graphs is 2-5% faster than FA2 graphs, so the graph gain is dispatch reduction, not the FA2 kernel. Async+graphs is 1-4% faster than FA2 graphs. The official GSM8K-87 runs agree in direction (eager 717.8 s, graphs 279.6 s, identical generated lengths).

Cost of the deterministic no-split schedule (kernel-only, one L20, page 16, random KV):

rows	ctx	split-KV ms	no split ms	ratio
1	4096	0.023	0.049	2.13x
1	8192	0.028	0.092	3.35x
8	4096	0.375	0.372	0.99x
8	8192	0.741	0.739	1.00x

At batch 1 and 8K context the no-split schedule adds roughly 6 ms per generated token (24 attention layers per depth view); at batch 8 it is unchanged. These are single-GPU, single-wave diagnostics, not production throughput claims.

Review

The maintainer review and the earlier audit are both addressed in the comment history. The per-layer CSR copy is reduced to one copy per traversal (one recorded group per capture), and the GSM8K --cuda-graphs non-native path now fails before reading the protocol or creating the output, with a CPU test. One smaller follow-up remains in the guide: a prime-time direct write, since a single copy per traversal is still needed because head counts are only known at the first call.

Limits

• Accuracy is exactly at the 59/87 baseline with determinism enabled.
• The CUDA-core FA2 path was bit-closer to the Torch reference but is not graph-safe on SM89; tensor cores are the only supported path here.
• FlashInfer JIT requires an nvcc >= 12.9 toolkit. Not validated on SM80/SM86, other page sizes/head dims, packed prefill, PD or speculation.
• Single GPU, single-wave timings; no production throughput claim.

Test plan

• CPU suite and FA2 GPU tests on one L20
• Real Ouro token parity across eager, graphs and async
• GSM8K-87 eager and graphs against the frozen 59/87 baseline
• Matched latency/throughput with Triton and FA2, eager and graphs

[sysuTobo]

Code audit of this draft (branch @ 86c09b7)

TL;DR: every traced data flow is correct, but the audit found one H-level defect — prepare_step synchronized the host on every traversal, contradicting the "device-only, no host read" claim and defeating async overlap. It is fixed, covered by a regression test, and re-validated on the L20. The fix is numerically neutral: GSM8K-87 reproduced 87/87 identical sequences before and after.

Audit target: this PR (feat/flashinfer-fa2-sm89 @ 86c09b7), base main. All items below were fixed unless marked open.

Findings

ID	Sev	Finding	Evidence	Status
H1	H	_write_csr assigned a Python scalar into a CUDA slice (csr.indptr[0] = 0), which stages through host memory and synchronizes. Under torch.cuda.set_sync_debug_mode("error"), prepare_step raised "called a synchronizing CUDA operation". Every traversal stalled the host once.	vllm_rlt/kernels/flashinfer_fa2_attention.py:214 (_write_csr)	Fixed: indptr[0].fill_(0), regression test test_prepare_step_does_not_synchronize; after the fix both prepare_step and __call__ report NO_SYNC
M1	M	Docs/PR claimed Triton graphs is 5-13% faster than FA2 graphs and a single "~3.2x/2.8x" graph factor. Measured values are 2-6% and 2.7-3.4x.	docs/flashinfer-fa2.md performance table; docs/accuracy.md matched-run paragraph	Fixed and refreshed with post-fix measurements
M2	M	The --cuda-graphs + non-native rejection ran after output.mkdir(), leaving an empty result directory on error and contradicting the ordering used by exit_settings.	benchmarks/gsm8k.py:590-596	Fixed: validation moved before mkdir
M3	M	Empty rows must report last_page_len == 1 (kernel contract), but tests only asserted the zeroed output; the CSR reference test had no empty row.	tests/test_flashinfer_fa2_attention.py:174-190	Fixed: empty row added and asserted
L1	L	No test for unsupported head ratios or for q/metadata row-count mismatch.	flashinfer_fa2_attention.py:239-243, 305-310	Fixed: two CPU tests
L2	L	Plan pinning/eviction (32-entry cache, capture-time pinning) has no test. Captured plans cannot be evicted by construction, but this is unverified.	flashinfer_fa2_attention.py:230-300	Open (follow-up)
L3	L	info["cuda_graphs"]/info["async_scheduling"] duplicate BackendCapabilities (two sources of truth).	flashinfer_fa2_attention.py:130-131	Open (diagnostic fields)
L4	L	__call__ does not validate q dtype/device, nor ceil(len/page) <= width; the engine guarantees both, but tampered device metadata could read past the CSR buffers.	flashinfer_fa2_attention.py:305-320	Open (contract note)
L5	L	The per-layer CSR copy scales with rows x table_width (rows=32, width=4096: about 0.5 MiB per layer, x24 layers x4 depth views, ~48 MiB per decode step).	flashinfer_fa2_attention.py:315-317	Open (follow-up: copy only the live prefix without a host read)

Data flows checked

• KV -> CSR -> FA2 -> output (_write_csr, __call__): real geometry Ouro-1.4B (16 heads, head_dim 128, page 16, table width 4096, rows <= max_num_seqs). indptr = cumsum(ceil(len/16)), compact indices via a fixed-shape permutation scatter, empty rows get last_page_len = 1, output rows with zero length are masked to zero. Device-only after the H1 fix.
• Priming -> capture -> replay (cuda_graph.py:289-304, 363-370, 469-491): priming happens after _copy_graph_inputs and before capture; replay re-primes before each launch; capture records the copy into the plan buffers, so replay consumes the refreshed staging. Verified by plan_misses == 0, one capture, 60 replays, and the adapter-level replay test with changed metadata.
• Plan/workspace ownership: one plan per (rows, q_heads, kv_heads, width), exact-sized workspace via workspace_size -> reset_workspace_buffer, capture-time pinning, eager plans evicted first.
• Capabilities -> execution routing: cuda_graphs/async_scheduling declared; device_context_lengths intentionally absent so async speculative is rejected at selection.
• GSM8K flag: --cuda-graphs -> metadata + Generator -> ExecutionConfig, fingerprint unchanged.

Coverage notes (searched, not inferred)

• grep -rn "prepare_attention" tests/ -> no direct calls; covered indirectly by engine tests and the plan_misses == 0 assertions.
• grep -rn "flashinfer_fa2" tests/ -> tests/test_flashinfer_fa2_attention.py (18 GPU + 5 CPU cases after this audit).
• grep -rn "cuda_graphs" tests/test_gsm8k*.py -> no hits: the harness flag itself is untested (L, recorded).
• Gate evidence: before the fix, the isolated _write_csr stage raised under sync-debug mode (the sync was indptr[0] assignment); after the fix both stages report NO_SYNC, so the new test fails on the old code. A deliberate mutation matrix was not run because the sandbox blocks mutated kernel executions.

Re-validation (one L20)

• CPU 24 passed (FA2 + GSM8K backends); GPU FA2 18 passed.
• Real Ouro-1.4B: Triton eager / FA2 eager / FA2 graphs / FA2 async+graphs produce identical tokens, plan_misses == 0.
• GSM8K-87: eager 60/87, graphs 59/87, both pass the 59/87 baseline; sequences identical pre/post fix for all 87 questions (eager 793.2 s, graphs 281.6 s).
• Post-fix matched timing (single session, C1/8/32): graphs are 2.7-3.4x over eager; Triton graphs is 2-6% faster than FA2 graphs; async+graphs is 1-4% faster than FA2 graphs.

Fix commits

• 6b5d2eb fix: keep FlashInfer FA2 CSR construction device-only
• 35f6d3b test: cover FA2 CSR bounds, row checks and synchronization
• 7a13bf9 bench: validate --cuda-graphs before creating the output
• 0edfdb7 docs: correct FA2 graph claims and refresh post-fix measurements

Open items

• L2 plan pinning/eviction test; L5 live-prefix CSR copy; the harness flag test noted above.

[0z5a]

Thanks for the careful scoping and the audit. I re-traced the backend contract (create_backend capability filtering, prepare_step priming in the eager / RecurrentGraphs / DeviceRecurrentGraphs / CodaGraphs paths, LAST_EXITED/SHARED tables, padding rows in bucketed graphs) and found no wrong-output bug. The vectorized CSR in _write_csr matches a reference construction on 2,000 random CPU cases (empty rows, -1 table padding, widths 1–40). The CPU suite passes on 86c09b7 (867 passed; GPU cases skipped). The performance write-up is honest: on the L20 table FA2 is never faster than Triton, so the value here is the #69 Ada coverage cell behind an explicit opt-in, which is fine.

Two things I would like addressed before this leaves draft, plus two smaller ones.

1. Eager plans are re-created as the context grows, and each new plan synchronizes the host — Medium

_wrapper keys plans on (rows, q_heads, kv_heads, width) (vllm_rlt/kernels/flashinfer_fa2_attention.py:232-233). On the eager path, width is the live table width, max(position // block_size + 1) (vllm_rlt/attention/metadata.py:64). It grows by one every 16 tokens of the longest row and changes with the prefill chunk size. So a steady eager decode creates a new plan at every page boundary. Every new plan:

• calls workspace_size() and plan(). Both do indptr.to("cpu") / last_page_len.to("cpu") on capacity_indptr / capacity_last, which are built on the device (:253-255), so each one blocks the host (FlashInfer 0.7.2.post1 flashinfer/decode.py, workspace_size and plan);
• calls reset_workspace_buffer(), which allocates a fresh pin_memory=True host buffer. The runner deliberately keeps cudaHostAlloc out of submission (vllm_rlt/worker/model_runner.py:241, :303);
• allocates new device buffers and a workspace.

I counted this on CPU with a fake wrapper around the PR's adapter (snippet below). One request with a 700-token prompt and 300 decode tokens, on the GSM8K harness shape (max_num_seqs=1, 128-row chunks), triggers 26 plan() calls. Eight staggered eager requests trigger 30. A working set wider than the 32-entry FIFO (:43, :288) re-plans on every step: 200 plans for 200 steps. The existing checks cannot catch this. plan_misses counts CSR priming misses, not plan creations. test_prepare_step_does_not_synchronize runs __call__ once, warm, before it enables sync-debug, and then checks only prepare_step. So "__call__ reports NO_SYNC" holds only while the plan cache is warm, and eager + async_scheduling loses its overlap at every page boundary. The module docstring (:5-6, "at the cache's full table capacity") describes the graph path only.

Suggested fix:

• Plan eager steps at the same fixed width the graph entries use (ceil(max_position_embeddings / block_size), as _graph_entry does), or bucket widths to powers of two, copying the live CSR into the prefix of the plan buffers. The key then depends effectively on rows only.
• Build capacity_indptr / capacity_last as CPU tensors so plan() / workspace_size() do no D2H.
• Expose a plan_creations counter next to plan_misses.
• Add a CPU test (fake wrapper, like the snippet) asserting a bounded plan count over a ~1K-token decode. This also covers the audit's open L2. Extend the sync-debug GPU test to a __call__ across a width change.

CPU reproduction (fake FlashInfer wrapper, PR adapter unchanged)

import sys, types, torch
calls = {"plan": 0}
class W:
    def __init__(self, *a, **k): pass
    def workspace_size(self, *a, **k): return 1024, 1024
    def reset_workspace_buffer(self, *a): pass
    def plan(self, *a, **k): calls["plan"] += 1
    def run(self, q, kv, out): out.zero_()
pkg = types.ModuleType("flashinfer"); pkg.__path__ = []
dec = types.ModuleType("flashinfer.decode"); dec.CUDAGraphBatchDecodeWithPagedKVCacheWrapper = W
sys.modules.update({"flashinfer": pkg, "flashinfer.decode": dec})
from vllm_rlt.kernels import flashinfer_fa2_attention as fa2
fa2.version = lambda n: "0.7.2.post1"; torch.version.hip = None
torch.cuda.get_device_capability = lambda d: (8, 9)
torch.cuda.is_current_stream_capturing = lambda: False
k = fa2.FlashInferFA2PagedAttention(torch.device("cuda"), torch.bfloat16, 128, 16)
k.device = torch.device("cpu")

def step(lengths):
    rows, width = len(lengths), max((l - 1) // 16 + 1 for l in lengths)  # metadata.py:64
    t = torch.randint(0, 1000, (rows, width), dtype=torch.int32)
    c = torch.tensor(lengths, dtype=torch.int32)
    k.prepare_step((types.SimpleNamespace(block_tables=t, context_lengths=c),))
    q = torch.zeros(rows, 16, 128, dtype=torch.bfloat16)
    kv = torch.zeros(1, 16, 16, 128, dtype=torch.bfloat16)
    k(q, kv, kv, t, c)

pos = 0
while pos < 700:                       # non-packed prefill: one row per token, 128-row chunks
    n = min(128, 700 - pos); step(list(range(pos + 1, pos + n + 1))); pos += n
for i in range(300):
    step([701 + i])
print(calls["plan"], k.plan_entries, k.plan_misses)   # -> 26 26 0

2. Eager and graph execution are not bit-identical; the split schedule depends on plan width and padded rows — Medium

The PR's own GSM8K-87 run shows FA2 eager and FA2 + CUDA Graphs disagree on 18/87 sequences, and the graph run lands exactly on the 59/87 gate. The docs attribute this to the full-width plan and padded graph rows changing the tensor-core split schedule. That is the expected behaviour of the default FA2 planner: the KV chunk size is chosen from the planned lengths and batch, so a row's reduction order depends on what it is planned with. The existing Triton backend is built to avoid exactly this (vllm_rlt/kernels/triton_attention.py:4-8, 15-16: "a row's partition depends only on its own length, so its result does not depend on the batch it runs in"). The #69 claim also made CUDA Graphs conditional on the graph path staying "bitwise-correct". Under adaptive exit, this variance also moves exit depths with batch composition and graph bucket.

The tiny-Ouro variant test cannot see it. _tiny_config() uses max_position_embeddings=128 (tests/test_flashinfer_fa2_attention.py:182), i.e. an 8-page plan, where no split happens. It is also below the 64+-page width at which the CUDA-core stall described in the docs appeared.

Suggested fix: pass disable_split_kv=True to both workspace_size() and plan(). FlashInfer documents this as the determinism control for CUDA Graph decode. With no split, a row's reduction order no longer depends on plan width or padded rows. That should make eager, graphs and any batch composition bit-identical, and it also makes the planned width irrelevant to the schedule, which simplifies finding 1. fixed_split_size is the alternative, but FlashInfer's docstring warns it is not guaranteed to be graph-compatible. The cost to check is low-batch, long-context throughput (fewer CTAs without splits). If that cost is unacceptable, please state the batch/graph variance as a documented limitation instead. Either way, please add a GPU test at a realistic width (≥ 64 pages, e.g. max_position_embeddings=2048): the same row alone vs inside a batch of 8, eager vs graph replay, compared bitwise (or with the expected tolerance if variance is kept).

3. Pin the FlashInfer kernel instead of relying on backend="auto" — Low

The wrapper is built without backend= (:257-258), i.e. "auto", while info hard-codes kernel="fa2" (:137). In 0.7.2.post1, auto resolves to FA2 for this graph wrapper (no caller-owned block_tables), but auto in 0.7.x can also route decode to cuDNN under some conditions. A patch release inside >=0.7.2,<0.8 could change the arithmetic silently. Fix: pass backend="fa2" explicitly and record wrapper.resolved_backend in info after the first plan().

4. Smaller items — Low

• Per-layer CSR copy (audit L5). Graph entries and plans are effectively 1:1 for a given (rows, width). prepare_step could write the CSR directly into the plan's indptr/indices/last_page_len (looked up at prime time) and drop the three per-layer copy_ calls at :322-324. For a 32-row bucket at width 4096 those copies move about 0.5 MiB per layer. Under graphs the recorded copy would then be unnecessary, since replay reads the plan buffers directly.
• Harness flag. benchmarks.gsm8k run --cuda-graphs / --attention-backend have no CPU test (the audit noted this too). A small argparse/metadata test, like the existing GSM8K backend tests, would keep the "not part of the fingerprint" claim from regressing.

[sysuTobo]

Thanks for the detailed review. All three main items are addressed in 86c09b7..d522a0d; the two smaller ones are recorded as follow-ups.

1. Eager plan churn — fixed

• Plan widths are rounded up to a power of two for the key, and the live CSR is copied into the prefix of the bucket-sized plan buffers.
• workspace_size() and plan() now receive pinned host capacity tensors, so plan creation no longer copies device metadata back to the host.
• plan_creations is exposed next to plan_misses.
• Your reproduction now reports 6 plans / 6 entries / 0 misses (was 26/26/0) for the same 700+300 workload; a CPU test asserts plan_creations <= 10 and plan_entries <= 10 over it.
• The sync-debug GPU test now warms two adjacent buckets (64 and 65 pages) and runs prepare_step + __call__ across them without synchronizing. Plan creation itself still happens once per bucket; the steady-state claim covers executions whose plans are warm.

2. Split-KV variance — deterministic path adopted

• disable_split_kv=True is passed to both workspace_size() and plan().
• Verified on the L20 with a 128-page plan and a 1,500-token row: with splitting, the row differs between a 128- and 256-page plan and alone vs inside a batch of eight; without splitting both comparisons are bitwise equal, and graph replay equals eager bitwise. New GPU tests cover width/batch/graph equality at 128 pages, and the tiny-Ouro engine test now uses max_position_embeddings=2048 so graph plans are 128 pages.
• GSM8K-87 after the change: eager 59/87, graphs 59/87, 87/87 identical sequences, no runaway (mean 106.7). The no-split arithmetic reproduces the earlier split-KV graph run on all 87 questions; the split-KV eager run (60/87) differed on 18.
• Measured cost (kernel-only, one L20, page 16): rows=1 ctx 4096/8192 → 2.13x/3.35x slower; rows=8 → 0.99x/1.00x. At 24 attention layers per depth view, batch 1 with 8K context adds roughly 6 ms/token; batch 8 is unchanged. Documented in the guide.

3. Pin the kernel — adapted to this FlashInfer version

plan() in 0.7.2.post1 has no backend argument and the graph wrapper's constructor does not expose one; the backend resolves inside _plan_impl. The adapter now requires wrapper.resolved_backend == "fa2" after the first plan (raising otherwise) and records it in info, and the selection info shares the adapter's dict so the value is visible after planning. This still fails loudly if a patch release routes to cuDNN.

4. Smaller items

• Per-layer CSR copy and the GSM8K flag test are recorded as follow-ups in the guide (not done in this round).
• The requested ≥64-page test is covered by the 128-page width/batch/graph test above.

Re-validation

• Full CPU suite: 896 passed / 7 skipped; FA2 + backend-selection GPU: 56 passed.
• Real Ouro-1.4B: Triton eager, FA2 eager, FA2 graphs and FA2 async+graphs remain token-identical (plan_misses == 0, one capture / 60 replays).
• Matched 64-token work (C1/8/32) with the final code: graphs 2.8-3.5x over eager; Triton graphs 2-5% faster than FA2 graphs; async+graphs 1-4% faster than FA2 graphs.

[sysuTobo]

Both smaller items are now in (9e170a4, 765fe41).

• Per-layer CSR copy (9e170a4). The staging-to-plan copy now runs once per traversal; later layers skip it, and a capture records one copy group instead of one per layer (previously 72 copies per traversal at 24 layers x 3 buffers). Replay still refreshes from the staging buffers written by prepare_step. Verified with the graph refresh/replay tests and real-model parity (1 capture / 60 replays, tokens identical to Triton).
• GSM8K flags (765fe41). run() now rejects --cuda-graphs with a non-native backend before reading the protocol or creating the output directory; test_cuda_graphs_requires_native_before_creating_output covers the failure and the absent output directory.

Full CPU suite: 897 passed / 7 skipped. FA2 GPU tests: 22 passed.

One residual follow-up is recorded in the guide: a prime-time direct write would remove the remaining single copy per traversal, but the plan key needs head counts that are only known at the first call.