[Feat][Attention] FA3 warp-specialized GQA prefill kernel (Hopper, fp16, causal)

benchmarks/benchmark_base.py

@@ -84,21 +84,52 @@ class BenchmarkWorkload(ShapeDtypeWorkload, InputGeneratingWorkload, Protocol):
 _bench_results = threading.local()
 
 
+# Name of the ``record_function`` annotation wrapping the timed call. Kineto
+# projects this scope onto the device timeline, so kernels the call launches
+# fall inside its window while the L2-flush ``cache.zero_()`` (enqueued outside
+# the scope) does not.
+_KERNEL_REGION = "tileops_bench_kernel"
+
+
 def _sum_kernel_time_us(kineto_results):
-    """Extract total CUDA kernel time directly from C++ Kineto events.
+    """Sum device time of the kernels the timed call launched.
+
+    Sums only kernels inside a :data:`_KERNEL_REGION` annotation window, so the
+    L2-flush fill is excluded and the kernel under test is counted regardless of
+    its name. A call launching several kernels contributes all of them.
+
+    Iterates the C++ Kineto events directly to bypass ``key_averages()``, which
+    is ~16x slower (~130ms of Python parsing/tree-building) for large traces.
 
-    Bypasses ``profiler.key_averages()`` which triggers expensive Python
-    event parsing (~120ms) and tree building (~10ms) for large traces.
-    Direct C++ iteration is ~16x faster for n_repeat=1280.
+    Returns:
+        ``(total_us, n_regions)``: summed kernel time in microseconds and the
+        number of annotation windows. The caller checks ``n_regions ==
+        n_repeat`` to confirm the scope projected on every iteration.
     """
-    total_us = 0.0
+    import bisect
+
+    windows: list[tuple[int, int]] = []
+    kernels: list[tuple[int, int]] = []  # (start_ns, duration_ns)
     for evt in kineto_results.events():
-        if evt.device_type() == DeviceType.CUDA:
-            name = evt.name()
-            if "vectorized_elementwise" in name and "FillFunctor" in name:
-                continue
-            total_us += evt.duration_ns() / 1000.0
-    return total_us
+        if evt.device_type() != DeviceType.CUDA:
+            continue
+        if evt.is_user_annotation():
+            if evt.name() == _KERNEL_REGION:
+                windows.append((evt.start_ns(), evt.end_ns()))
+            continue
+        kernels.append((evt.start_ns(), evt.duration_ns()))
+
+    windows.sort()
+    starts = [w[0] for w in windows]
+    ends = [w[1] for w in windows]
+    total_us = 0.0
+    for start_ns, dur_ns in kernels:
+        # Count only kernels that fall inside a timed-call window; everything
+        # outside (notably the L2-flush fill) is excluded.
+        idx = bisect.bisect_right(starts, start_ns) - 1
+        if idx >= 0 and start_ns < ends[idx]:
+            total_us += dur_ns / 1000.0
+    return total_us, len(windows)
 
 
 # ---------------------------------------------------------------------------
@@ -198,38 +229,46 @@ def _run(i):
         _run(i % n_repeat)
     torch.cuda.synchronize()
 
-    # Timed trials with CUPTI (single profiler, n_trials cycles)
+    # Timed trials with CUPTI.  Each trial opens its own torch.profiler context
+    # around exactly n_repeat iterations and reads the trace after the context
+    # closes; summed device kernel time / n_repeat is the mean per-call kernel
+    # time.  We deliberately do NOT use torch.profiler.schedule: that mechanism
+    # is for sampling a window out of a long step()-driven loop, and forcing
+    # n_repeat calls into a single "step" let queued, un-synchronized launches
+    # leak across the warmup/active boundary.  A plain per-trial context records
+    # exactly the calls we want — no schedule, no on_trace_ready callback.
+    #
+    # Only the timed call is wrapped in record_function(_KERNEL_REGION), so the
+    # L2-flush cache.zero_() enqueued just before it is attributed by scope (not
+    # kernel name) and excluded.  Flush and call share the stream, so the flush
+    # completes before the call begins (L2 cold) without a sync between them; we
+    # sync only after the call so its kernels are recorded before the next flush.
     trial_means: list[float] = []
-
-    def _on_trace_ready(prof):
-        kr = prof.profiler.kineto_results
-        kernel_us = _sum_kernel_time_us(kr) / n_repeat
-        trial_means.append(kernel_us * 1e-3)
-
     try:
         with suppress_stdout_stderr():
-            schedule = torch.profiler.schedule(
-                wait=0, warmup=1, active=1, repeat=n_trials,
-            )
-            profiler = torch.profiler.profile(
-                activities=[torch.profiler.ProfilerActivity.CUDA],
-                schedule=schedule,
-                on_trace_ready=_on_trace_ready,
-            )
-            with profiler:
-                for _ in range(n_trials):
-                    # Warmup step (discarded by schedule)
-                    for i in range(n_repeat):
-                        cache.zero_()
-                        _run(i)
-                    profiler.step()
-                    # Active step (measured → _on_trace_ready)
+            for _ in range(n_trials):
+                with torch.profiler.profile(
+                    # CPU activity is required for Kineto to project the
+                    # annotation onto the device timeline (CUDA-only emits no
+                    # window); it adds only host-side overhead, not kernel time.
+                    activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.CUDA,
+                    ],
+                ) as profiler:
                     for i in range(n_repeat):
                         cache.zero_()
-                        _run(i)
-                    profiler.step()
+                        with torch.profiler.record_function(_KERNEL_REGION):
+                            _run(i)
+                        torch.cuda.synchronize()  # kernel recorded in isolation
+                total_us, n_regions = _sum_kernel_time_us(profiler.profiler.kineto_results)
+                # Scope failed to project on some iteration → trace untrustworthy,
+                # fall back to CUDA events.
+                if n_regions != n_repeat:
+                    raise RuntimeError
+                trial_means.append((total_us / n_repeat) * 1e-3)
     except RuntimeError:
-        pass
+        trial_means = []
 
     # Fallback to CUDA events if CUPTI failed
     if not trial_means:

benchmarks/ops/attention/bench_gqa.py

@@ -100,22 +100,25 @@ def baseline_fn(q, k, v, o, grad_output, lse):
     return baseline_fn
 
 
-def _flashinfer_gqa_fwd(test: GroupedQueryAttentionFwdTest, q, k, v):
-    """Set up FlashInfer batched prefill wrapper. Returns callable or None."""
+def _flashinfer_gqa_fwd(test, q, k, v):
+    """FlashInfer ragged-prefill baseline. Handles seq_len_q != seq_len_kv (square is
+    the seq_len_q == seq_len_kv case). Returns callable or None."""
     try:
         from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper  # noqa: PLC0415
     except ImportError:
         return None
 
-    B, S, H, D = q.shape
+    B, Sq, H, D = q.shape
+    Skv = k.shape[1]
     Hkv = k.shape[2]
-    cu_seqlens = torch.arange(0, B + 1, dtype=torch.int32, device=q.device) * S
+    qo_indptr = torch.arange(0, B + 1, dtype=torch.int32, device=q.device) * Sq
+    kv_indptr = torch.arange(0, B + 1, dtype=torch.int32, device=q.device) * Skv
 
     workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
     wrapper = BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD")
     wrapper.plan(
-        qo_indptr=cu_seqlens,
-        kv_indptr=cu_seqlens,
+        qo_indptr=qo_indptr,
+        kv_indptr=kv_indptr,
         num_qo_heads=H,
         num_kv_heads=Hkv,
         head_dim_qk=D,
@@ -128,7 +131,7 @@ def run_fn(q, k, v):
             q.reshape(-1, H, D),
             k.reshape(-1, Hkv, D),
             v.reshape(-1, Hkv, D),
-        ).reshape(B, S, H, D)
+        ).reshape(B, Sq, H, D)
 
     return run_fn
 
@@ -393,6 +396,11 @@ def test_gqa_prefill_fwd_bench(
     result_bl = bm.profile(_torch_gqa_prefill_ref(test), *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
+    fi_fn = _flashinfer_gqa_fwd(test, *inputs)
+    if fi_fn is not None:
+        result_fi = bm.profile(fi_fn, *inputs)
+        BenchmarkReport.record(op, locals(), result_fi, tag="flashinfer")
+
 
 _GQA_PREFILL_VARLEN_FWD_BENCH_PARAMS = [
     pytest.param(

tileops/kernels/attention/__init__.py

@@ -30,6 +30,7 @@
 )
 from .gqa_fwd_fp8 import GQAFwdFP8Fa3ContractPtxAccBN224WsTmaVKernel
 from .gqa_fwd_ws import GQAFwdWsPersistentCausalKernel, GQAFwdWsPersistentKernel
+from .gqa_prefill_fwd_ws import GQAPrefillFwdWsPersistentCausalKernel
 from .gqa_prefill_varlen_fwd import GQAPrefillVarlenFwdKernel
 from .gqa_sliding_window_fwd import (
     GQASlidingWindowFwdKernel,
@@ -55,6 +56,7 @@
     "GQAFwdWsPersistentCausalKernel",
     "GQAFwdWsPersistentKernel",
     "GQAPrefillFwdKernel",
+    "GQAPrefillFwdWsPersistentCausalKernel",
     "GQAPrefillPagedWithFP8KVCacheFwdKernel",
     "GQAPrefillPagedWithKVCacheFwdKernel",
     "GQAPrefillPagedWithKVCacheRopeAppendKernel",