Support CUDA-graph capture on AITER batch decode (opt-in)

CHANGELOG.md

@@ -5,6 +5,7 @@
 - Add Jupyter notebook tutorial for using amd-flashinfer on ROCm (#213) @diptorupd
 - Add ROCm profiler module (`flashinfer/rocm_profiler/`) and FA2 single-prefill benchmark driver using rocprofv3 (#205) @diptorupd
 - Gate `torch.compile` integration behind `FLASHINFER_USE_TORCH_CUSTOM_OPS`, with HIP pytest coverage (#210) @demandal25
+- Support CUDA-graph capture on the AITER batch-decode path via an explicit `backend="aiter"`: capture at a maximum sequence length and replay for shorter sequences (the launch grid and `.so` variant are fixed at capture-time shapes; the kernel early-exits per sequence on `context_lens`). Previously `backend="aiter"` raised under `use_cuda_graph=True`. `backend="auto"` continues to use the in-tree `fa2` kernel under capture (fa2's graph path is capacity-based and capture-order-independent). Adds `benchmarks/rocm_benchmarks/bench_decode_graph_hip.py` and HIP graph-replay tests. @demandal25
 
 ## Changed
 

README.md

@@ -55,7 +55,7 @@ kernel for non-attention ops). **AITER** = ROCm AITER backend.
 | Kernel | HIP | AITER | `backend="auto"` resolves to | Notes |
 | :--- | :---: | :---: | :--- | :--- |
 | **Single decode attention** | ✅ `fa2` | — | HIP | MHA / GQA / MQA |
-| **Batch decode attention (paged)** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + `pos_encoding_mode="NONE"` + no CUDA-graph + `use_tensor_cores=False`; else **HIP** | MHA / GQA / MQA; **fp8 KV-cache (E4M3FNUZ)** on the HIP path; sliding-window on the AITER path; CUDA-graph auto-routes back to HIP |
+| **Batch decode attention (paged)** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + `pos_encoding_mode="NONE"` + no CUDA-graph + `use_tensor_cores=False`; else **HIP** | MHA / GQA / MQA; **fp8 KV-cache (E4M3FNUZ)** on the HIP path; sliding-window on the AITER path; CUDA-graph capture on the AITER path is **opt-in via `backend="aiter"`** (grid + `.so` fixed at capture-time shapes — capture at max seq len); `backend="auto"` uses HIP `fa2` under capture |
 | **Single prefill attention** | ✅ `fa2` | ✅ | **AITER** when `fp16/bf16` + `NHD` + no custom mask + equal Q/KV dtypes & head dims + `pos_encoding_mode="NONE"`; else **HIP** | MHA / GQA / MQA; fp8 WIP |
 | **Batch prefill attention (paged + ragged)** | ✅ `fa2` | ✅ | Same auto criteria as single prefill | MHA / GQA / MQA; fp8 WIP. AITER native page sizes: `{16, 1024}` (`{128, 256, 1024}` on `amd-aiter==0.1.10`); other sizes go through a gather on the AITER path |
 | **Cascade attention** | ✅ | — | HIP | Two-level shared-prefix attention; a fused single-kernel HIP variant is gated behind `FLASHINFER_HIP_FUSED_CASCADE=1` |
@@ -354,9 +354,14 @@ Backend-specific exceptions to "auto picks AITER when supported":
 * `rmsnorm`: `backend="auto"` picks the AITER C++ path (CK `rmsnorm2d`) only
   for 2-D fp16/bf16 inputs whose weight dtype matches; 3-D inputs, fp32, or a
   mismatched weight dtype fall back to the HIP `native` kernel.
-* `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True`
-  force `auto` back to `fa2` (AITER decode does not support either),
-  and `pos_encoding_mode != "NONE"` raises under `backend="aiter"`.
+* `batch_decode`: `use_cuda_graph=True` or `use_tensor_cores=True` force `auto`
+  back to `fa2`, and `pos_encoding_mode != "NONE"` raises under
+  `backend="aiter"`. CUDA-graph capture on the AITER decode path is available
+  via an explicit `backend="aiter"` (not `auto`): capture at your maximum
+  sequence length — the grid and `.so` variant are fixed at capture-time shapes
+  and the kernel early-exits per sequence, so replays *shorter* than captured
+  are correct but *longer* ones are not. fa2's graph path is capacity-based and
+  carries no such constraint.
 
 Unless you are using the prebuilt Docker image, install AITER separately
 via one of the options below.
@@ -396,7 +401,7 @@ you need any of them: `backend="fa2"` for attention wrappers, or
 * `q_dtype != kv_dtype` (mixed-precision Q/KV is unsupported)
 * `head_dim_qk != head_dim_vo` (e.g. DeepSeek-style MLA with 192/128 head dims)
 * `pos_encoding_mode != "NONE"` (AITER attention paths only support `"NONE"`)
-* batch decode: `use_cuda_graph=True` or `use_tensor_cores=True`
+* batch decode: `use_tensor_cores=True`
 * the `aiter` Python package is not importable
 
 **Features silently ignored on the AITER path** (kwargs are accepted by

benchmarks/rocm_benchmarks/bench_batch_decode_hip.py

@@ -0,0 +1,205 @@
+"""
+Copyright (c) 2026 Advanced Micro Devices, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Batch paged-decode benchmark: fa2 (in-tree HIP) vs aiter (PA v1), eager.
+
+Baseline for the "HIP-graph-safe fast decode" work item
+(docs/rocm_library_optimization_plan.md §3 #1): today AITER decode falls back
+to fa2 under CUDA-graph capture, so the fast path never runs in the serving
+hot loop. This bench measures the eager fa2-vs-aiter gap (the win we forfeit
+under graphs) across a serving-shaped batch x kv_len sweep. Once the
+capacity-sized-grid change lands, add a --graph path here to compare
+aiter-under-graph vs fa2-under-graph and to measure the idle-block
+over-provisioning cost (the one open perf item in the plan).
+
+Shapes: decode (q_len=1), GQA 32/8, HD=128, bf16, page_size=16.
+Sweep: batch x kv_len.
+
+Run:
+    python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py               # full pipeline
+    python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --timing-only # no profiling
+    python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --backend fa2
+    python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --backend aiter
+    python benchmarks/rocm_benchmarks/bench_batch_decode_hip.py --counters stall
+
+Design note: bench flags are parsed at module level because rocprofv3
+re-executes this script as a subprocess per PMC pass with the same sys.argv.
+Module-level parsing ensures the subprocess builds identical configs to the
+outer timing run.
+"""
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+import torch
+
+import flashinfer
+from flashinfer.jit.core import logger as _jit_logger
+
+# Suppress routine JIT INFO/DEBUG output; WARNING still surfaces compile errors.
+_jit_logger.setLevel(logging.WARNING)
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "rocm_profiler"))
+from rocm_profiler import KernelConfig, RocmProfiler
+
+# ---------------------------------------------------------------------------
+# Bench-script-level argument parsing (see design note above)
+# ---------------------------------------------------------------------------
+_bench_parser = argparse.ArgumentParser(add_help=False)
+_bench_parser.add_argument(
+    "--counters",
+    default="roofline",
+    metavar="PRESET_OR_FILE",
+    help=(
+        "Counter preset name ('roofline', 'occupancy', 'stall', 'compute', "
+        "'memory', 'basic') or path to a rocprofv3 YAML file. Default: roofline."
+    ),
+)
+_bench_parser.add_argument(
+    "--label",
+    default=None,
+    metavar="PREFIX",
+    help="Output-file label prefix (default: 'decode' for roofline, 'decode_<preset>' otherwise).",
+)
+_bench_parser.add_argument(
+    "--backend",
+    default="both",
+    choices=["fa2", "aiter", "both"],
+    help="Which decode backend(s) to sweep. Default: both.",
+)
+_bench_args, _ = _bench_parser.parse_known_args()
+
+_counters = _bench_args.counters
+_label = (
+    _bench_args.label
+    if _bench_args.label is not None
+    else ("decode" if _counters == "roofline" else f"decode_{_counters}")
+)
+
+# ---------------------------------------------------------------------------
+# Sweep configuration — decode is q_len=1; shapes ~ Llama-70B TP1 decode.
+# ---------------------------------------------------------------------------
+_NUM_QO_HEADS = 32
+_NUM_KV_HEADS = 8
+_HEAD_DIM = 128
+_DTYPE = torch.bfloat16
+_PAGE_SIZE = 16
+_BATCHES = [1, 8, 32, 128, 256]
+_KV_LENS = [1024, 2048, 4096, 8192]
+
+_OUTPUT_DIR = str(Path(__file__).parent)
+
+
+def _flops(kv_len: int, num_qo_heads: int, head_dim: int) -> int:
+    # decode: q_len=1 → attended = kv_len; QK^T + PV ≈ 2 matmuls → factor 4.
+    return kv_len * num_qo_heads * head_dim * 4
+
+
+def _bytes(kv_len: int, num_qo_heads: int, num_kv_heads: int, head_dim: int) -> int:
+    # Dominated by reading K and V (q_len=1). 2 bytes/elem (bf16).
+    return 2 * head_dim * (2 * num_qo_heads + 2 * kv_len * num_kv_heads)
+
+
+def _build_paged_kv(batch, kv_len, page_size, num_kv_heads, head_dim, dtype, device):
+    """Build batch-paged decode KV + query tensors for one (batch, kv_len)."""
+    num_full_pages, last_tokens = divmod(kv_len, page_size)
+    if last_tokens == 0:
+        last_tokens = page_size
+    else:
+        num_full_pages += 1
+    total_pages = num_full_pages * batch
+
+    kv_data = torch.randn(
+        total_pages, 2, page_size, num_kv_heads, head_dim, dtype=dtype, device=device
+    )
+    kv_indptr = (
+        torch.arange(batch + 1, dtype=torch.int32, device=device) * num_full_pages
+    )
+    _rng = torch.Generator(device=device).manual_seed(42)
+    kv_indices = torch.randperm(
+        total_pages, dtype=torch.int32, device=device, generator=_rng
+    )
+    kv_last_page_len = torch.full(
+        (batch,), last_tokens, dtype=torch.int32, device=device
+    )
+    q = torch.randn(batch, _NUM_QO_HEADS, head_dim, dtype=dtype, device=device)
+    return q, kv_data, kv_indptr, kv_indices, kv_last_page_len
+
+
+def _make_wrapper_config(backend, batch, kv_len):
+    device = torch.device("cuda")
+    q, kv_data, kv_indptr, kv_indices, kv_last_page_len = _build_paged_kv(
+        batch, kv_len, _PAGE_SIZE, _NUM_KV_HEADS, _HEAD_DIM, _DTYPE, device
+    )
+    ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=device)
+    wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(ws, "NHD", backend=backend)
+    wrapper.plan(
+        kv_indptr,
+        kv_indices,
+        kv_last_page_len,
+        _NUM_QO_HEADS,
+        _NUM_KV_HEADS,
+        _HEAD_DIM,
+        _PAGE_SIZE,
+        pos_encoding_mode="NONE",
+        q_data_type=_DTYPE,
+        kv_data_type=_DTYPE,
+    )
+    return KernelConfig(
+        name=f"{backend}_b{batch}_kv{kv_len}",
+        run_fn=torch.inference_mode()(lambda q=q, kv=kv_data, w=wrapper: w.run(q, kv)),
+        theoretical_flops=batch * _flops(kv_len, _NUM_QO_HEADS, _HEAD_DIM),
+        theoretical_bytes=batch
+        * _bytes(kv_len, _NUM_QO_HEADS, _NUM_KV_HEADS, _HEAD_DIM),
+        num_tokens=batch,  # one decoded token per sequence
+        label=f"{backend:<5s} b={batch:>3d}  kv={kv_len:>5d}",
+    )
+
+
+@torch.inference_mode()
+def _make_configs() -> list[KernelConfig]:
+    backends = (
+        ["fa2", "aiter"] if _bench_args.backend == "both" else [_bench_args.backend]
+    )
+    if "aiter" in backends and not flashinfer.aiter_utils.is_aiter_supported(
+        torch.device("cuda:0")
+    ):
+        print("AITER not supported on this device; dropping the aiter sweep.")
+        backends = [b for b in backends if b != "aiter"]
+
+    configs: list[KernelConfig] = []
+    for backend in backends:
+        for batch in _BATCHES:
+            for kv_len in _KV_LENS:
+                configs.append(_make_wrapper_config(backend, batch, kv_len))
+    return configs
+
+
+if __name__ == "__main__":
+    _skip_gpu = "--replot" in sys.argv or "--list-presets" in sys.argv
+    profiler = RocmProfiler(
+        configs=[] if _skip_gpu else _make_configs(),
+        num_warmup=3,
+        dry_run_ms=100,
+        repeat_ms=1000,
+        counters=_counters,
+        kernel_name_regex="",
+        output_dir=_OUTPUT_DIR,
+        label=_label,
+        roofline=(_counters == "roofline"),
+    )
+    profiler.run()