feat: add FlashInfer attention backend for SM120 decode

README.md

@@ -56,9 +56,9 @@ vLLM to be installed. See [Citation](#citation) for the paper's BibTeX entry.
 - **Depth-aware paged KV cache.** LAST-EXITED and SHARED layouts, automatic
   CUDA cache sizing, and optional prefix caching, incremental page allocation,
   priority scheduling, and preemption with CPU state snapshots.
-- **Configurable GPU execution.** Triton and FlashAttention backends, asynchronous
-  scheduling, multiple CUDA streams, reusable buffers, and CUDA Graph capture
-  of decode recurrent cores.
+- **Configurable GPU execution.** Triton, FlashAttention, and FlashInfer backends,
+  asynchronous scheduling, multiple CUDA streams, reusable buffers, and CUDA Graph
+  capture of decode recurrent cores.
 - **Offline and online inference.** A Python API, a command-line interface,
   and an OpenAI-compatible completions endpoint with streaming, greedy decoding,
   and seeded top-k/top-p sampling.

benchmarks/gsm8k.py

@@ -439,6 +439,7 @@ def run(args):
     output.mkdir(parents=True, exist_ok=False)
     metadata = {
         "backend": args.backend,
+        "attention_backend": args.attention_backend,
         "protocol_sha256": file_digest(args.protocol),
         "source": source(),
         "cuda_visible_devices": os.environ["CUDA_VISIBLE_DEVICES"],
@@ -456,7 +457,10 @@ def run(args):
     tokenizer = AutoTokenizer.from_pretrained(model, local_files_only=True)
     task = make_task(protocol["task_config"])
     start = time.monotonic()
-    generator = Generator(args.backend, str(model), tokenizer, protocol["max_length"], exit_policy)
+    generator = Generator(
+        args.backend, str(model), tokenizer, protocol["max_length"], exit_policy,
+        attention_backend=args.attention_backend,
+    )
     load_seconds = time.monotonic() - start
     rows = []
     with (output / "samples.jsonl").open("x", buffering=1) as stream:
@@ -619,6 +623,12 @@ def main():
     p.add_argument("--backend", choices=["transformers", "native"], required=True)
     p.add_argument("--protocol", required=True)
     p.add_argument("--output", required=True)
+    p.add_argument(
+        "--attention-backend",
+        choices=["torch", "triton", "flashinfer"],
+        default="triton",
+        help="Native engine attention kernel; recorded in the run metadata",
+    )
     p.add_argument("--exit-mode", choices=["ouro", "ouro_delayed"], default="ouro")
     p.add_argument(
         "--exit-threshold",

benchmarks/gsm8k_backends.py

@@ -80,7 +80,9 @@ def _after_model(self, module, inputs, output):
 
 
 class Generator:
-    def __init__(self, backend, model_path, tokenizer, max_length, exit_policy):
+    def __init__(
+        self, backend, model_path, tokenizer, max_length, exit_policy, attention_backend="triton"
+    ):
         self.backend, self.tokenizer, self.exit = backend, tokenizer, exit_policy
         # Full depth, taken from the protocol (the checkpoint's total_ut_steps).
         loops = exit_policy["max_loops"]
@@ -113,7 +115,7 @@ def __init__(self, backend, model_path, tokenizer, max_length, exit_policy):
                 model_path,
                 device="cuda",
                 dtype=torch.bfloat16,
-                attention_backend="triton",
+                attention_backend=attention_backend,
                 # last_exited KV keeps one plane per loop.
                 cache_config=CacheConfig(num_blocks=loops * ((max_length + 15) // 16)),
                 scheduler_config=SchedulerConfig(max_num_seqs=1, max_num_batched_tokens=128),

docs/README.md

@@ -16,6 +16,7 @@ an environment to receiving the first generated response. The
 | [Runtime configuration](cdb_runtime.md) | Exit policies, KV layouts, execution options, and CUDA Graphs |
 | [Asynchronous scheduling](https://github.com/hsliuustc0106/vllm-rlt/pull/30) | CPU/GPU pipelining and single-stream or multi-stream execution |
 | [FlashAttention](https://github.com/hsliuustc0106/vllm-rlt/pull/30) | FA2/FA3/FA4 installation, hardware selection, and constraints |
+| [FlashInfer](flashinfer.md) | XQA decode kernels for SM120: installation, environment, and scope |
 | [Cache and scheduling features](https://github.com/hsliuustc0106/vllm-rlt/pull/31) | Prefix reuse, incremental KV, priorities, and preemption |
 | [Prefill/decode disaggregation](https://github.com/hsliuustc0106/vllm-rlt/pull/31) | Single-host GPU worker pools and NIXL transfer |
 | [HTTP serving](serving.md) | Completions API, streaming, and service lifecycle |
@@ -26,7 +27,8 @@ an environment to receiving the first generated response. The
 
 Current model support is limited to Ouro-1.4B. CPU execution provides a Torch
 reference backend. FlashAttention hardware validation is currently documented
-for FA4 on B300; FA2/FA3 require validation on their target devices. Disaggregated
+for FA4 on B300; FA2/FA3 require validation on their target devices. FlashInfer
+decode is validated on RTX 5090 (SM120). Disaggregated
 serving currently targets multiple GPUs on a single host.
 
 `ouro_delayed` reuses the trained Ouro gate with a one-loop delay; it changes

docs/accuracy.md

@@ -164,6 +164,25 @@ counts as `identical_context` when the two runs chose different depths for a pos
 whose preceding tokens were still identical. These fields are diagnostic; the accuracy
 gate does not use them.
 
+## Attention-backend runs
+
+`run --backend native` selects the attention kernel with `--attention-backend`
+(`triton` by default, recorded in the run metadata):
+
+```bash
+python -m benchmarks.gsm8k run --backend native --attention-backend flashinfer \
+  --protocol /path/to/gsm8k-87-protocol.json --output /path/to/native-flashinfer
+```
+
+On RTX 5090 (SM120) the frozen default case scored **Triton 61/87 and FlashInfer
+58/87** against the stored HF baseline of 59/87: 64 of the 87 generations were
+token-identical between the kernels, the divergent ones flipped near-tied greedy
+choices at scattered positions, and exit depths never differed. This is the
+expected single-pass greedy variation between BF16 kernels with different
+accumulation orders — see the [FlashInfer accuracy
+note](flashinfer.md#accuracy-note-gsm8k-87-single-greedy-pass) before treating a
+one-question delta from a backend swap as a regression.
+
 Dataset: [GSM8K](https://huggingface.co/datasets/openai/gsm8k).
 The [Ouro evaluation settings](https://arxiv.org/html/2510.25741v5#A3.T16) do not
 pin the exact harness revision, demonstrations or token limits, so the settings

docs/flashinfer.md

@@ -0,0 +1,97 @@
+# FlashInfer Attention Backend
+
+The `flashinfer` backend runs paged decode attention through FlashInfer's XQA
+kernels (`flashinfer-python >= 0.7.0`, `trtllm_batch_decode_with_kv_cache`,
+`backend="auto"`). It is the validated attention path for NVIDIA consumer
+Blackwell (SM120, e.g. RTX 5090), where the pinned FlashAttention releases
+reject paged KV.
+
+## Scope
+
+- Decode and per-token prefill, identical semantics to the Triton path,
+  including CUDA Graph capture and asynchronous scheduling.
+- No packed prefill: the engine keeps the FA4-only packed path closed, so
+  prefill runs per token exactly as with Triton.
+- No prefill/decode disaggregation; PD entrypoints reject the backend
+  explicitly.
+- `bfloat16` only. `float16` is rejected at startup: the kernels compute
+  correctly, but flashinfer-python 0.7.0.post1 crashes the process (host
+  SIGFPE) while building the f16 JIT module on first in-process use.
+- `--block-size` must be one of `16, 32, 64, 128` (compiled XQA page-size
+  variants) and `head_dim` must be divisible by 16 within `[16, 512]`.
+
+## Installation and runtime requirements
+
+```bash
+pip install -e ".[flashinfer]"
+```
+
+The XQA kernels are compiled just in time on first use, which imposes two
+environment requirements beyond the package install:
+
+1. **A CUDA toolkit with `nvcc >= 12.9` reachable via `CUDA_HOME`.** SM120
+   kernels are built for the `compute_120f` architecture family, which needs
+   CUDA 12.9 or newer. Without `nvcc`, startup falls back to comparing against
+   the CUDA version PyTorch was built with, and SM 12.x devices are rejected
+   with `SM 12.x requires CUDA >= 12.9`.
+2. **That toolkit's `lib` directory on `LD_LIBRARY_PATH`.** The compiled
+   modules link `libcudart.so.13`; pip-installed toolkits keep it outside the
+   default loader path. With the pip toolkit layout
+   (`nvidia-cuda-nvcc`/`nvidia-cuda-runtime`/`nvidia-cuda-cccl`/...), the
+   unified tree lives under `site-packages/nvidia/cu13` and may need
+   `lib64/libcudart.so` and `lib64/stubs/libcuda.so` developer symlinks, which
+   the wheels do not ship.
+
+The first call on a machine compiles for a few seconds and caches under
+`~/.cache/flashinfer/<package-version>/<arch>/`; later calls, including CUDA
+Graph capture warmups, reuse the cache.
+
+## Running
+
+```bash
+vllm-rlt --model <ouro-checkpoint> --revision <revision> \
+    --device cuda --dtype bfloat16 --attention-backend flashinfer
+```
+
+The adapter holds one 128 MiB zero-initialized workspace for the kernel's
+scratch and softmax statistics; it is allocated once at startup and reused, so
+graph capture performs no new allocations. Padded block-table columns are
+clamped non-negative before the call: the engine produces `-1` padding on the
+synchronous path, `0` on the asynchronous path, and stale-but-valid ids in
+CUDA Graph replay beyond the current width. The kernel bounds reads by
+`seq_lens`, so all three padding states stay invisible.
+
+## Choosing between Triton and FlashInfer
+
+SM120 is the only validated target for this backend; on it, FlashInfer decodes
+noticeably faster at low concurrency (single-request decode measured ~2x
+against the Triton kernel at 256-token contexts) because the XQA kernel
+parallelizes a single query row across KV chunks, while the Triton kernel
+spawns one program per (row, head). At high concurrency with long contexts the
+two are comparable; choose Triton when its determinism properties matter more
+and FlashInfer for single-user or low-batch latency. SM9/SM10 devices are
+accepted but unvalidated — prefer FlashAttention there until validated
+configurations are recorded.
+
+## Validation status
+
+Validated on RTX 5090 (SM120, 32 GB) with Ouro-1.4B BF16, `block_size=16`,
+fixed four-loop decode: oracle parity against the reference paged attention,
+token/exit-depth equality against the Triton backend across synchronous,
+asynchronous, and multi-stream execution with static buffers and graph replay,
+and end-to-end greedy generation.
+
+## Accuracy note (GSM8K-87, single greedy pass)
+
+On the frozen GSM8K-87 protocol (HF baseline 59/87), a fixed four-loop greedy
+pass scores 58/87 with FlashInfer versus 61/87 with Triton. 64 of 87
+generations are token-identical between the two backends; the remaining 23
+diverge at scattered positions with unchanged exit depths, and the net
+correct-answer flips are 3 against, 0 for. Kernel-level comparison against a
+float32 oracle at 352-1024 token contexts shows the Triton kernel is
+near-exact (max deviation up to 5e-4) while XQA's split reductions deviate by
+up to 4e-3 — within normal BF16 kernel variation, but enough to flip
+near-tied greedy argmax decisions during long chain-of-thought decoding.
+Triton remains the default backend; treat this note as the expected greedy
+variation when opting into FlashInfer on SM120.
+

pyproject.toml

@@ -26,6 +26,7 @@ pd = ["nixl==1.4.1"]
 serve = []
 flash-attn = ["flash-attn>=2.8.3"]
 flash-attn4 = ["flash-attn-4==4.0.0b30"]
+flashinfer = ["flashinfer-python>=0.7.0"]
 triton = []
 dev = ["pytest>=8", "ruff>=0.9"]
 

tests/test_backend_selection.py

@@ -15,9 +15,15 @@
 
 @pytest.fixture
 def packages(monkeypatch):
-    from vllm_rlt.kernels import flash_attention
-
-    installed = {"flash_attn", "flash_attn_3.flash_attn_interface", "flash_attn.cute", "triton"}
+    from vllm_rlt.kernels import flash_attention, flashinfer_attention
+
+    installed = {
+        "flash_attn",
+        "flash_attn_3.flash_attn_interface",
+        "flash_attn.cute",
+        "triton",
+        "flashinfer",
+    }
     original_import = flash_attention.importlib.import_module
 
     def import_module(name):
@@ -32,10 +38,17 @@ def import_module(name):
                 flash_attn_with_kvcache=lambda *args, **kwargs: None,
                 flash_attn_varlen_func=lambda *args, **kwargs: None,
             )
+        if name == "flashinfer.decode":
+            if "flashinfer" not in installed:
+                raise ImportError("flashinfer is not installed")
+            return SimpleNamespace(
+                trtllm_batch_decode_with_kv_cache=lambda *args, **kwargs: None
+            )
         return original_import(name)
 
     monkeypatch.setattr(flash_attention.importlib, "import_module", import_module)
     monkeypatch.setattr(flash_attention, "version", lambda name: "test")
+    monkeypatch.setattr(flashinfer_attention, "version", lambda name: "test")
     monkeypatch.setattr(torch.version, "hip", None)
     return installed
 
@@ -74,7 +87,8 @@ def test_cpu_auto_reports_torch_and_runs_reference_attention(caplog):
             "flash_attn_2",
         ),
         ((10, 0), torch.bfloat16, 128, 16, (), "flash_attn_4"),
-        ((12, 0), torch.bfloat16, 128, 16, (), "triton"),
+        ((12, 0), torch.bfloat16, 128, 16, (), "flashinfer"),
+        ((12, 0), torch.bfloat16, 128, 16, ("flashinfer",), "triton"),
         ((7, 5), torch.float16, 128, 256, (), "torch"),
     ],
 )
@@ -109,12 +123,15 @@ def test_auto_reports_missing_packages_and_configuration(packages, monkeypatch):
     assert "multiple of 256" in reason and "triton" in reason
 
 
-@pytest.mark.parametrize("backend", ["flash_attn_2", "triton"])
-def test_explicit_missing_package_never_falls_back(packages, monkeypatch, backend):
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9))
+@pytest.mark.parametrize(
+    "backend,sm,block_size",
+    [("flash_attn_2", (8, 9), 256), ("triton", (8, 9), 256), ("flashinfer", (12, 0), 16)],
+)
+def test_explicit_missing_package_never_falls_back(packages, monkeypatch, backend, sm, block_size):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: sm)
     packages.clear()
     with pytest.raises(ImportError):
-        create_backend(backend, torch.device("cuda"), torch.bfloat16, 128, 256)
+        create_backend(backend, torch.device("cuda"), torch.bfloat16, 128, block_size)
 
 
 def test_explicit_flash_incompatibility_is_not_overridden(packages, monkeypatch):