[bench] vit_attention: model-matching V stride + controllable cache-eviction (--flush-mib)

benchmarks/vit_attention/bench.py

@@ -20,6 +20,24 @@
 Triton-specific tuning knobs (--bm/--bn/--nw/--ns/--we) are passed directly
 to the kernel launch and are ignored for other backends.
 
+Operand layout: a real ViT fuses the QKV projection, so V is a *non-contiguous*
+view into a packed tensor (token stride = mult*H*D, mult=3 for a standard fused
+QKV) while Q/K are made contiguous by the serving wrappers. --v-stride-mult
+controls this (default 3 = model-matching fused QKV; 1 = fully contiguous).
+
+Cache working-set state is controllable, because it dominates the gap vs real
+serving (where surrounding qkv-proj / MLP / norm work only *partially* evicts
+q/k/v between attention calls):
+  * default            -> stock do_bench (fully cold; clears the L2 each rep)
+  * --no-flush         -> fully warm (q/k/v stay resident between reps)
+  * --flush-mib N      -> zero an N-MiB device buffer before each rep, streaming
+                          it through the last-level cache to displace q/k/v.
+                          N=0 is warm; N ~ LLC size models serving-like partial
+                          eviction; N >= 2x LLC is fully cold. (A write equal to
+                          the cache size does not fully evict due to set/way
+                          associativity + the flush stream evicting itself, so
+                          full-cold needs ~2x the cache size.)
+
 Output: one JSON line per backend to stdout.
 """
 
@@ -86,11 +104,89 @@ def _parse() -> argparse.Namespace:
     p.add_argument(
         "--we", type=int, default=None, help="waves_per_eu (TRITON_ATTN only)"
     )
+    p.add_argument(
+        "--v-stride-mult",
+        type=int,
+        default=3,
+        help=(
+            "V token-stride multiplier: V is a non-contiguous slice of a packed "
+            "(B,S,mult,H,D) tensor, so its token stride is mult*H*D. Default 3 "
+            "matches a fused QKV projection (the real model layout); 1 = the "
+            "old fully-contiguous layout."
+        ),
+    )
+    p.add_argument(
+        "--no-flush",
+        action="store_true",
+        help=(
+            "Measure with a warm cache (no L2 flush between iterations) instead "
+            "of do_bench's default cold-cache flush. Equivalent to --flush-mib 0."
+        ),
+    )
+    p.add_argument(
+        "--flush-mib",
+        type=int,
+        default=None,
+        help=(
+            "Zero an N-MiB device buffer before each timed rep to model a "
+            "controllable amount of cache eviction. N=0 = warm (no flush); "
+            "N ~ last-level-cache size = serving-like partial eviction; "
+            "N >= 2x LLC = fully cold. Overrides --no-flush. When unset, the "
+            "default stock do_bench full-cold L2 flush is used."
+        ),
+    )
     p.add_argument("--warmup-ms", type=int, default=200)
     p.add_argument("--rep-ms", type=int, default=600)
     return p.parse_args()
 
 
+def _do_bench_flush(fn, warmup_ms: int, rep_ms: int, flush_mib: int) -> list[float]:
+    """do_bench variant with a controllable N-MiB cache flush between reps.
+
+    Mirrors triton.testing.do_bench's event-timed loop (5-call cost estimate ->
+    warmup -> timed reps). Before each timed rep it zeroes a ``flush_mib``-MiB
+    device buffer; those writes stream through the last-level cache and displace
+    the resident q/k/v. ``flush_mib=0`` performs no flush (fully warm), so q/k/v
+    stay resident; ``flush_mib ~ LLC`` models serving-like partial eviction and
+    ``flush_mib >= 2x LLC`` is fully cold. Returns per-call times in ms.
+    """
+    di = triton.runtime.driver.active.get_device_interface()
+
+    flush_buf = None
+    if flush_mib > 0:
+        flush_buf = torch.empty(
+            int(flush_mib) * 1024 * 1024, dtype=torch.int8, device="cuda"
+        )
+
+    fn()
+    di.synchronize()
+
+    start = di.Event(enable_timing=True)
+    end = di.Event(enable_timing=True)
+    start.record()
+    for _ in range(5):
+        fn()
+    end.record()
+    di.synchronize()
+    estimate_ms = start.elapsed_time(end) / 5
+    n_warmup = max(1, int(warmup_ms / estimate_ms))
+    n_repeat = max(1, int(rep_ms / estimate_ms))
+
+    for _ in range(n_warmup):
+        fn()
+    starts = [di.Event(enable_timing=True) for _ in range(n_repeat)]
+    ends = [di.Event(enable_timing=True) for _ in range(n_repeat)]
+    di.synchronize()
+    for i in range(n_repeat):
+        if flush_buf is not None:
+            flush_buf.zero_()
+        starts[i].record()
+        fn()
+        ends[i].record()
+    di.synchronize()
+    return [s.elapsed_time(e) for s, e in zip(starts, ends)]
+
+
 def _bench_one(
     backend_name: str,
     args: argparse.Namespace,
@@ -209,13 +305,26 @@ def _fn():
     else:
         raise ValueError(f"Unsupported backend: {backend_name}")
 
-    # do_bench clears the L2 cache before each measurement iteration.
-    all_times = triton.testing.do_bench(
-        _fn,
-        warmup=args.warmup_ms,
-        rep=args.rep_ms,
-        return_mode="all",
-    )
+    if args.flush_mib is not None:
+        # Controllable partial eviction: zero an N-MiB buffer before each rep.
+        all_times = _do_bench_flush(_fn, args.warmup_ms, args.rep_ms, args.flush_mib)
+        flush_mode = "warm" if args.flush_mib == 0 else "flush_mib"
+        flush_mib = args.flush_mib
+    elif args.no_flush:
+        # Warm cache: no L2 flush between iterations (== --flush-mib 0).
+        all_times = _do_bench_flush(_fn, args.warmup_ms, args.rep_ms, 0)
+        flush_mode = "warm"
+        flush_mib = 0
+    else:
+        # do_bench clears the L2 cache before each measurement iteration.
+        all_times = triton.testing.do_bench(
+            _fn,
+            warmup=args.warmup_ms,
+            rep=args.rep_ms,
+            return_mode="all",
+        )
+        flush_mode = "cold"
+        flush_mib = None
     all_times = sorted(all_times)
     n = len(all_times)
     mean = sum(all_times) / n
@@ -231,6 +340,12 @@ def _fn():
         "D": D,
         "dtype": args.dtype,
         "backend": backend_name,
+        "v_stride_mult": args.v_stride_mult,
+        "v_token_stride": int(v4.stride(-3)),
+        "v_contiguous": bool(v4.is_contiguous()),
+        "flush": flush_mode != "warm",
+        "flush_mode": flush_mode,
+        "flush_mib": flush_mib,
     }
     if backend == AttentionBackendEnum.TRITON_ATTN:
         config.update(
@@ -271,9 +386,22 @@ def main() -> int:
 
     # Wrappers (FLASH_ATTN, ROCM_AITER_FA, TRITON_ATTN) expect (B, S, H, D).
     # TORCH_SDPA expects the same via vit_torch_sdpa_wrapper.
+    #
+    # Q/K are contiguous (the serving wrappers make them so). V models the fused
+    # QKV projection: it is a non-contiguous slice of a packed (B,S,mult,H,D)
+    # tensor, giving a token stride of mult*H*D (mult=3 => the real fused-QKV
+    # layout). --v-stride-mult 1 restores the old fully-contiguous V.
     q4 = torch.randn(B, S, H, D, dtype=dtype, device=device)
     k4 = torch.randn(B, S, H, D, dtype=dtype, device=device)
-    v4 = torch.randn(B, S, H, D, dtype=dtype, device=device)
+    mult = args.v_stride_mult
+    if mult <= 1:
+        v4 = torch.randn(B, S, H, D, dtype=dtype, device=device)
+    else:
+        # Packed buffer; take index 0 along the fused dim -> V keeps a token
+        # stride of mult*H*D. The (B,S) dims stay mergeable, so the TRITON path's
+        # .view(B*S, H, D) still works while preserving the non-contiguous stride.
+        v_packed = torch.randn(B, S, mult, H, D, dtype=dtype, device=device)
+        v4 = v_packed[:, :, 0]
 
     # cu_seqlens / max_seqlen used by FA and Triton backends
     cu = torch.tensor([i * S for i in range(B + 1)], dtype=torch.int32, device=device)