diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/floor.py b/benchmarks/kernels/rocm_rdna_hip_attn/tools/floor.py new file mode 100644 index 000000000000..956318b5a3d0 --- /dev/null +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/floor.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The most a decode-attention kernel can reach against utils' roofline. + +The roofline counts one dispatch plus the bytes at peak bandwidth. A paged +kernel also pays a dependent page-table read and DRAM latency, which at short +context are most of the time. This times a kernel that only streams the same +bytes after one dependent page-table read, cold (rotated through a >= 96 MiB +working set) under a HIP graph, as the attention harness does. Its %roof at a +configuration's bytes is that configuration's ceiling. + + PYTHONPATH=$PWD .venv/bin/python \\ + benchmarks/kernels/rocm_rdna_hip_attn/tools/floor.py \\ + --hq 16 --hkv 16 --head-dim 32 --m 1 4 + +Prints the stream floor per size, then the ceiling per configuration, M and +context. +""" + +import argparse +import bisect +import math +import pathlib +import sys + +import torch +from torch.utils.cpp_extension import load_inline + +from vllm.triton_utils import triton + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +import utils # noqa: E402 + +_SRC = r""" +#include +#include +#include +typedef unsigned int u4 __attribute__((ext_vector_type(4))); +// Each workgroup reads its page index, the dependent load a paged kernel +// cannot avoid, then streams its share of the buffer in b128. +template +__global__ __launch_bounds__(256) void stream(const u4* __restrict__ p, + const int* __restrict__ bt, + long n16, unsigned* out) { + const long stride = (long)gridDim.x * 256 * UNR; + u4 acc = {0, 0, 0, 0}; + const long base = __builtin_amdgcn_readfirstlane(bt[blockIdx.x & 7]); + for (long i = (long)blockIdx.x * 256 * UNR + threadIdx.x; i < n16; + i += stride) { + u4 v[UNR]; +#pragma unroll + for (int u = 0; u < UNR; ++u) { + const long j = i + (long)u * 256 + base; + v[u] = j < n16 ? p[j] : u4{0, 0, 0, 0}; + } +#pragma unroll + for (int u = 0; u < UNR; ++u) acc ^= v[u]; + } + const unsigned r = acc.x ^ acc.y ^ acc.z ^ acc.w; + if (r == 0x12345678u) out[blockIdx.x] = r; +} +__global__ void empty(unsigned* out) { + if (threadIdx.x == 999) out[0] = 1; +} +void run(torch::Tensor buf, torch::Tensor bt, int64_t n16, int64_t grid, + int64_t unr, torch::Tensor out) { + auto s = at::cuda::getCurrentCUDAStream(); + auto* p = (const u4*)buf.data_ptr(); + auto* o = (unsigned*)out.data_ptr(); + const int* b = bt.data_ptr(); + if (unr < 0) { + hipLaunchKernelGGL(empty, dim3(1), dim3(32), 0, s, o); + return; + } +#define L(U) hipLaunchKernelGGL((stream), dim3(grid), dim3(256), 0, s, p, b, n16, o) + if (unr == 1) L(1); else if (unr == 2) L(2); else if (unr == 4) L(4); else L(8); +} +""" + +CONTEXTS = [128, 512, 1024, 4096, 8192, 16384, 32768] + + +def _module(): + sdk = pathlib.Path(torch.__file__).resolve().parents[1] / "_rocm_sdk_devel" + return load_inline( + "rdna_floor_probe", + cpp_sources="void run(torch::Tensor, torch::Tensor, int64_t, int64_t, " + "int64_t, torch::Tensor);", + cuda_sources=_SRC, + functions=["run"], + extra_cuda_cflags=["-O3"], + extra_ldflags=[f"-L{sdk / 'lib'}"], + ) + + +def _bench(mod, bufs, bts, n16, grid, unr, out): + def fn(): + for b, t in zip(bufs, bts): + mod.run(b, t, n16, grid, unr, out) + + fn() + torch.accelerator.synchronize() + ms = triton.testing.do_bench_cudagraph(fn, return_mode="median") + return ms * 1000 / len(bufs) + + +def measure(mod, sizes_kib, working_set_mb=96): + """Best-of-grid stream time, in us, for each size.""" + out = torch.zeros(4096, dtype=torch.int32, device="cuda") + floor = {} + for kib in sizes_kib: + nbytes = kib * 1024 + layers = max(10, -(-working_set_mb * 1024**2 // nbytes)) + bufs = [ + torch.randint(0, 255, (nbytes,), dtype=torch.uint8, device="cuda") + for _ in range(layers) + ] + bts = [torch.zeros(8, dtype=torch.int32, device="cuda") for _ in bufs] + n16 = nbytes // 16 + best = math.inf + for grid in (8, 16, 20, 40, 80, 160, 320, 640, 1280): + for unr in (1, 2, 4, 8): + if grid * 256 * unr > 4 * n16 and grid > 8: + continue + best = min(best, _bench(mod, bufs, bts, n16, grid, unr, out)) + floor[kib] = best + del bufs, bts + torch.accelerator.empty_cache() + return floor + + +def floor_us(floor, nbytes): + """Log-log interpolation of the measured floor at nbytes.""" + pts = sorted(floor.items()) + xs = [math.log(k) for k, _ in pts] + x = math.log(max(nbytes / 1024, pts[0][0])) + i = min(max(bisect.bisect(xs, x), 1), len(xs) - 1) + (_, t0), (_, t1) = pts[i - 1], pts[i] + f = (x - xs[i - 1]) / (xs[i] - xs[i - 1]) + return math.exp(math.log(t0) + f * (math.log(t1) - math.log(t0))) + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--hq", type=int, required=True) + p.add_argument("--hkv", type=int, required=True) + p.add_argument("--head-dim", type=int, required=True) + p.add_argument("--window", type=int, default=0) + p.add_argument("--m", type=int, nargs="+", default=[1, 4]) + p.add_argument("--contexts", type=int, nargs="+", default=CONTEXTS) + p.add_argument( + "--sizes", + type=int, + nargs="+", + default=[ + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 65536, + 262144, + ], + help="stream sizes in KiB; they must bracket every configuration's KV", + ) + args = p.parse_args() + + mod = _module() + out = torch.zeros(16, dtype=torch.int32, device="cuda") + tiny = [torch.zeros(16, dtype=torch.uint8, device="cuda") for _ in range(100)] + bts = [torch.zeros(8, dtype=torch.int32, device="cuda") for _ in tiny] + print(f"empty kernel: {_bench(mod, tiny, bts, 1, 1, -1, out):.2f} us") + + floor = measure(mod, args.sizes) + print("\n| KiB | floor us | roofline us | %roof |") + print("| --- | --- | --- | --- |") + for kib, us in floor.items(): + roof = utils.DISPATCH_US + kib * 1024 / (utils.PEAK_GIBS * 1024**3) * 1e6 + print(f"| {kib} | {us:.2f} | {roof:.2f} | {roof / us * 100:.1f} % |") + + hq, hkv, d, w = args.hq, args.hkv, args.head_dim, args.window + print(f"\n{hq}/{hkv}/{d}" + (f" w{w}" if w else "") + ": ceiling (%roof)") + print("| M | " + " | ".join(str(s) for s in args.contexts) + " |") + print("| --- " * (len(args.contexts) + 1) + "|") + for m in args.m: + r = [] + for s in args.contexts: + roof = utils.roofline_us(hq, hkv, d, m, s, window=w) + nbytes = (roof - utils.DISPATCH_US) * 1e-6 * utils.PEAK_GIBS * 1024**3 + r.append(roof / floor_us(floor, nbytes) * 100) + print(f"| {m} | " + " | ".join(f"{x:.1f} %" for x in r) + " |") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/jit.py b/benchmarks/kernels/rocm_rdna_hip_attn/tools/jit.py index 482a87da8050..3e67f3cc9ee1 100644 --- a/benchmarks/kernels/rocm_rdna_hip_attn/tools/jit.py +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/jit.py @@ -130,6 +130,10 @@ class _PrefillArgs(ctypes.Structure): ("num_entries", ctypes.c_int), ("bt_stride", ctypes.c_int), ("head_group", ctypes.c_int), + ("v", ctypes.c_void_p), + ("q_stride", ctypes.c_int), + ("k_stride", ctypes.c_int), + ("v_stride", ctypes.c_int), ] @@ -374,6 +378,55 @@ def prefill_attn( ) self._launch(ctypes.byref(args)) + def encoder_attn( + self, + q, + k, + v, + out, + partial_o, + partial_ml, + counters, + seq_lens, + max_seq_len, + scale, + plan=None, + query_start_loc=None, + head_group=0, + ): + var = self.variant + if var.causal: + raise RuntimeError("a causal build runs through prefill_attn") + for x in (q, k, v): + if x.stride(2) != 1 or x.stride(1) != var.head_size: + raise RuntimeError("each token's heads must be contiguous") + if not (out.is_contiguous() and q.dtype == k.dtype == v.dtype == var.dtype): + raise RuntimeError("out must be contiguous, all of the build's dtype") + if (plan is None) != (query_start_loc is None): + raise RuntimeError("a batch takes both a plan and query_start_loc") + args = _PrefillArgs( + q=q.data_ptr(), + kv=k.data_ptr(), + out=out.data_ptr(), + seq_lens=seq_lens.data_ptr(), + partial_o=partial_o.data_ptr(), + partial_ml=partial_ml.data_ptr(), + counters=counters.data_ptr(), + num_tokens=q.size(0), + max_seq_len=max_seq_len, + scale=scale, + stream=torch.cuda.current_stream(q.device).cuda_stream, + plan=plan.data_ptr() if plan is not None else None, + query_start=query_start_loc.data_ptr() if plan is not None else None, + num_entries=plan.size(0) if plan is not None else 0, + head_group=head_group, + v=v.data_ptr(), + q_stride=q.stride(0), + k_stride=k.stride(0), + v_stride=v.stride(0), + ) + self._launch(ctypes.byref(args)) + def load(v: R.KernelVariant | R.PrefillVariant) -> Jit | JitPrefill | None: """The variant's build, or None if it fails (see `failure`). diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/matrix.py b/benchmarks/kernels/rocm_rdna_hip_attn/tools/matrix.py index 20352f995e8e..0d44257e442a 100644 --- a/benchmarks/kernels/rocm_rdna_hip_attn/tools/matrix.py +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/matrix.py @@ -7,7 +7,9 @@ `path` is what ran (`kernel`, `prefill` or `triton`); %roof is against utils.roofline. --batch takes batch specs or SPECS names; --baselines adds backends to compare. Forcing a knob, --ablate or --jit builds the variants -with jit.py. +with jit.py. --non-causal times the mm-encoder path against TRITON_ATTN's, +TORCH_SDPA's and, with flash-attn installed, FLASH_ATTN's on segments of +L = S tokens, each through the wrapper MMEncoderAttention calls. PYTHONPATH=$PWD python \\ benchmarks/kernels/rocm_rdna_hip_attn/tools/matrix.py @@ -78,8 +80,12 @@ def load(args: argparse.Namespace) -> tuple[list[Shape], int]: """ shapes: list[Shape] = [] windowed = 0 + # Decoder shapes run causal, encoder shapes non-causal (--non-causal). + mode = "encoder" if getattr(args, "non_causal", False) else "decoder" with open(args.shapes) as fh: for r in csv.DictReader(fh): + if not int(r.get(mode, mode == "decoder")): + continue s = Shape(int(r["Hq"]), int(r["Hkv"]), int(r["D"]), int(r["window"])) if args.hq and s.hq not in args.hq: continue @@ -202,6 +208,147 @@ def summary(baselines: list[str], ratios: dict[str, list[float]]) -> None: ) +def _segments(spec: str) -> list[int]: + """Segments of a spec: 577x5 is five of 577, 1024,257 those two.""" + if "x" in spec: + n, count = spec.split("x") + return [int(n)] * int(count) + return [int(x) for x in spec.split(",")] + + +def _encoder_runs(hq, hkv, d, segs, dtype): + """The baselines' runs by backend name and ours, over 10 layers of + inputs, and the first Q. Each through the wrapper MMEncoderAttention + calls, given what the encoder gives it: cu_seqlens on the device, its + host lengths noted as maybe_recompute_cu_seqlens does, max_seqlen on the + host. FLASH_ATTN only with flash-attn installed.""" + import numpy as np + import torch + + from vllm.v1.attention.backends import fa_utils + from vllm.v1.attention.ops.vit_attn_wrappers import ( + flash_attn_maxseqlen_wrapper, + rdna_attn_wrapper, + remember_rdna_cu_seqlens, + torch_sdpa_wrapper, + triton_attn_wrapper, + ) + + n = sum(segs) + layers = [ + [ + (torch.randn(1, n, h, d, device="cuda") * 0.5).to(dtype) + for h in (hq, hkv, hkv) + ] + for _ in range(10) + ] + host = np.concatenate([[0], np.cumsum(segs)]).astype(np.int32) + cu_seqlens = torch.from_numpy(host).cuda() + remember_rdna_cu_seqlens(cu_seqlens, host) + max_seqlen = torch.tensor(max(segs), dtype=torch.int32) + scale = d**-0.5 + + def triton(): + for q, k, v in layers: + triton_attn_wrapper(q, k, v, 1, scale, cu_seqlens, max_seqlen) + + def sdpa(): + for q, k, v in layers: + torch_sdpa_wrapper(q, k, v, scale, cu_seqlens, enable_gqa=hq != hkv) + + def flash(): + for q, k, v in layers: + flash_attn_maxseqlen_wrapper( + q, k, v, 1, False, None, scale, cu_seqlens, max_seqlen + ) + + def ours(): + for q, k, v in layers: + rdna_attn_wrapper(q, k, v, 1, scale, cu_seqlens, max_seqlen) + + runs = {"TRITON_ATTN": triton, "TORCH_SDPA": sdpa} + if fa_utils._ROCM_FLASH_ATTN_AVAILABLE: + runs["FLASH_ATTN"] = flash + # The cu_seqlens tensor lives with the runs, so its lengths stay noted. + return runs, ours, layers[0][0][0] + + +def non_causal(args: argparse.Namespace, shapes: list[Shape]) -> None: + """--non-causal: per shape and segment list, both mm-encoder kernels.""" + import torch + + from vllm.v1.attention.ops import rocm_rdna_encoder + + dtype = utils.torch_dtype(args.dtype) + specs = args.segments or [str(m) for m in args.m] + + def timed(runs, layers=10): + """Median us per layer of each run, timed in turn.""" + for run in runs: + for _ in range(3): + run() + ts: list[list[float]] = [[] for _ in runs] + for _ in range(20 * args.reps): + for run, t in zip(runs, ts): + a = torch.cuda.Event(enable_timing=True) + b = torch.cuda.Event(enable_timing=True) + a.record() + run() + b.record() + b.synchronize() + t.append(a.elapsed_time(b) * 1e3 / layers) + return [statistics.median(t) for t in ts] + + from vllm.v1.attention.backends import fa_utils + + baselines = ["TRITON_ATTN", "TORCH_SDPA"] + if fa_utils._ROCM_FLASH_ATTN_AVAILABLE: + baselines.append("FLASH_ATTN") + print(f"") + columns = ["segments", "Hq", "Hkv", "D", "bound", "roofline", *baselines] + columns += ["ours", *(f"x {b}" for b in baselines), "%roof"] + columns += [f"{b} %roof" for b in baselines] + print("| " + " | ".join(columns) + " |") + print("| --- " * len(columns) + "|") + ratios: dict[str, list[float]] = {b: [] for b in baselines} + for spec in specs: + segs = _segments(spec) + for hq, hkv, d, window in shapes: + if window: + continue + roof, bound = utils.roofline_noncausal( + hq, hkv, d, segs, torch.finfo(dtype).bits // 8 + ) + runs, ours, q = _encoder_runs(hq, hkv, d, segs, dtype) + why = rocm_rdna_encoder.unsupported(q, hkv) + base_runs = [runs[b] for b in baselines] + served = [] + for run in base_runs: + try: + run() + served.append(run) + except Exception: # noqa: BLE001 - e.g. SDPA's math path OOMs + torch.accelerator.empty_cache() + times = iter(timed(served if why else [*served, ours])) + base = [next(times) if run in served else None for run in base_runs] + head = [spec, hq, hkv, d, bound, f"{roof:.2f}"] + head += ["-" if t is None else f"{t:.2f}" for t in base] + if why: + print(f"| {' | '.join(map(str, head))} | {why} |", flush=True) + continue + us = next(times) + for b, t in zip(baselines, base): + if t is not None: + ratios[b].append(t / us) + cells = [*head, f"{us:.2f}"] + cells += ["-" if t is None else f"{t / us:.2f}x" for t in base] + cells += [ + "-" if t is None else f"{roof / t * 100:.1f} %" for t in (us, *base) + ] + print("| " + " | ".join(map(str, cells)) + " |", flush=True) + summary(baselines, ratios) + + def main() -> None: p = argparse.ArgumentParser() add_arguments(p) @@ -219,6 +366,16 @@ def main() -> None: help="one sequence of --m new tokens after each of --prefixes cached ones", ) p.add_argument("--prefixes", type=int, nargs="+", default=[0, 4096, 16384]) + p.add_argument( + "--non-causal", + action="store_true", + help="the mm-encoder path: segments of --m tokens, or --segments", + ) + p.add_argument( + "--segments", + nargs="+", + help="with --non-causal, one call's segments: 577x5 or 1024,257,9", + ) p.add_argument( "--contexts", type=int, @@ -285,7 +442,15 @@ def main() -> None: from runner import run_attention_benchmark if args.m is None: - args.m = [128, 256, 512, 1024, 2048, 4096, 8192] if args.prefill else [1, 4] + args.m = ( + [128, 256, 512, 1024, 2048, 4096, 8192] + if args.prefill or args.non_causal + else [1, 4] + ) + if args.non_causal: + shapes, _ = load(args) + non_causal(args, shapes) + return if args.batch: specs = [s for b in args.batch for s in SPECS.get(b, [b])] elif args.prefill: diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/shapes.csv b/benchmarks/kernels/rocm_rdna_hip_attn/tools/shapes.csv index c1dc6671c7e5..330e3d441d99 100644 --- a/benchmarks/kernels/rocm_rdna_hip_attn/tools/shapes.csv +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/shapes.csv @@ -1,34 +1,55 @@ -Hq,Hkv,D,window -14,2,64,0 -16,2,64,0 -32,8,64,0 -32,32,64,0 -32,32,96,0 -10,10,128,0 -16,2,128,0 -16,8,128,0 -24,8,128,0 -28,4,128,0 -32,2,128,0 -32,4,128,0 -32,8,128,0 -32,32,128,0 -40,8,128,0 -8,1,256,0 -8,2,256,0 -8,4,256,0 -16,2,256,0 -16,4,256,0 -16,8,256,0 -24,4,256,0 -8,1,512,0 -8,2,512,0 -16,1,512,0 -16,2,512,0 -32,4,512,0 -8,1,256,512 -8,2,256,512 -8,4,256,1024 -8,4,256,4096 -16,8,256,1024 -32,16,256,1024 +Hq,Hkv,D,window,decoder,encoder +4,4,32,0,0,1 +8,8,32,0,0,1 +16,16,32,0,0,1 +32,32,32,0,0,1 +32,32,48,0,0,1 +16,16,72,0,0,1 +16,16,80,0,0,1 +14,2,64,0,1,0 +16,2,64,0,1,0 +32,8,64,0,1,0 +32,32,64,0,1,0 +4,4,64,0,0,1 +6,6,64,0,0,1 +8,8,64,0,0,1 +12,12,64,0,0,1 +14,14,64,0,0,1 +16,16,64,0,0,1 +20,20,64,0,0,1 +24,24,64,0,0,1 +64,64,64,0,0,1 +32,32,96,0,1,0 +10,10,128,0,1,0 +16,2,128,0,1,0 +16,8,128,0,1,0 +24,8,128,0,1,0 +28,4,128,0,1,0 +32,2,128,0,1,0 +32,4,128,0,1,0 +32,8,128,0,1,0 +32,32,128,0,1,0 +40,8,128,0,1,0 +4,4,128,0,0,1 +8,8,128,0,0,1 +12,12,128,0,0,1 +16,16,128,0,0,1 +24,24,128,0,0,1 +8,1,256,0,1,0 +8,2,256,0,1,0 +8,4,256,0,1,0 +16,2,256,0,1,0 +16,4,256,0,1,0 +16,8,256,0,1,0 +24,4,256,0,1,0 +8,1,512,0,1,0 +8,2,512,0,1,0 +16,1,512,0,1,0 +16,2,512,0,1,0 +32,4,512,0,1,0 +8,1,256,512,1,0 +8,2,256,512,1,0 +8,4,256,1024,1,0 +8,4,256,4096,1,0 +16,8,256,1024,1,0 +32,16,256,1024,1,0 diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/tune.py b/benchmarks/kernels/rocm_rdna_hip_attn/tools/tune.py index bdc374412530..c3eaf5a2958f 100644 --- a/benchmarks/kernels/rocm_rdna_hip_attn/tools/tune.py +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/tune.py @@ -102,8 +102,23 @@ def grid(hq, hkv, d, m, batch=False): if max_segments == 1 and min_segment_blocks != 1: continue for head_dim_split in sorted({rule, 2 * rule}): - lane_dims = (d // head_dim_split) // 16 - if lane_dims not in (4, 6, 8) or waves % head_dim_split: + lane_dims = -(-d // 16) // head_dim_split + if ( + lane_dims - lane_dims % 2 not in (2, 4, 6, 8) + or waves % head_dim_split + ): + continue + # As the kernel's asserts: whole chunks per part, an odd lane + # slice unsplit, the score exchange inside a wave's K tile. + row_tiles = -(-(rows // row_groups) // 16) + if ( + -(-d // 16) % head_dim_split + or (lane_dims % 2 and head_dim_split != 1) + or ( + head_dim_split > 1 + and row_tiles * 1024 > 16 * ((d // head_dim_split) * 2 + 16) + ) + ): continue out.append( { @@ -120,7 +135,7 @@ def grid(hq, hkv, d, m, batch=False): } ) # The dot path needs whole dot2 pairs per lane (not D=96). - if not batch and (8 if d == 128 else d // 32) % 2 == 0: + if not batch and (8 if d == 128 else max(2, d // 32)) % 2 == 0: for waves, max_segments in itertools.product(*DOT_GRID.values()): out.append({"dot_product": 1, "waves": waves, "max_segments": max_segments}) return out @@ -729,6 +744,9 @@ def prefill_grid(d): if ( k["waves"] % (split * k["key_waves"]) or k["key_tile"] // 16 % k["key_waves"] + # An odd lane slice (D=48, 80) or a padded one (D=72) needs V + # transposed. + or ((-(-d // 16) // split % 2 or d % 16) and not k["v_t"]) ): continue out.append(k) @@ -1122,6 +1140,310 @@ def pid(knobs): fh.write(json.dumps(dec) + "\n") +# Non-causal (an encoder's segments, L = S): the sequence lengths tuned, the +# steps' own and the traced mm-encoder calls'. +NON_CAUSAL_CELLS = tuple( + sorted({*STEPS, 320, 577, 784, 1200, 2040, 3200, 3520, 4292, 8192}) +) +# Finer than STEPS: one row per size band of the encoders' calls. +NON_CAUSAL_STEPS = (*STEPS[:6], 768, 1024, 1536, 2048, 3072, 4096, 8192) +NON_CAUSAL_GRID = { + "waves": (4, 8, 16), + "key_tile": (16, 32, 64), + "v_t": (0, 1), + "key_waves": (1, 2), + "double_buffer": (0, 1), + "max_segments": (1, 4, 16), + "head_group": (1, 2, 4, 8), + "row_tiles": (1, 2), +} + + +def non_causal_grid(d): + """Non-causal knob sets for head size d: the whole space, as the masking + it drops changes the register pressure.""" + split = max(1, d // 128) + out = [] + for values in itertools.product(*NON_CAUSAL_GRID.values()): + k = dict(zip(NON_CAUSAL_GRID, values), head_dim_split=split) + if ( + k["waves"] % (split * k["key_waves"]) + or k["key_tile"] // 16 % k["key_waves"] + or ((-(-d // 16) // split % 2 or d % 16) and not k["v_t"]) + or ( + k["row_tiles"] > 1 + and ( + split > 1 + or k["key_waves"] > 1 + or k["key_tile"] > 32 + # compute2 reads V^T. + or not k["v_t"] + ) + ) + ): + continue + out.append(k) + return out + + +def tune_prefill_non_causal(a) -> None: + """Prefill rows for one sequence of L = S tokens, every token seeing every + key: per step of M, the knobs fastest against TRITON_ATTN's mm-encoder + kernel (context_attention_fwd, is_causal=False) over NON_CAUSAL_CELLS.""" + import jit + import torch + from triton.testing import do_bench + + import vllm + import vllm.v1.attention.ops.rocm_rdna_hip as R + from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd + + jit.install() + dtype = utils.torch_dtype(a.dtype) + hq, hkv, d, bs = a.hq, a.hkv, a.head_dim, a.block_size + if a.window: + raise SystemExit("a non-causal build has no window") + tag = f"{hq}/{hkv}/{d} non-causal prefill {a.dtype}" + cells = [m for m in NON_CAUSAL_CELLS if m >= STEPS[0]] + flops = [4.0 * hq * d * m * m for m in cells] + weight = {c: math.sqrt(f / sum(flops)) for c, f in zip(cells, flops)} + roof = {m: utils.roofline_noncausal(hq, hkv, d, [m])[0] for m in cells} + tol = 5e-3 if dtype == torch.float16 else 3e-2 + fixed = { + "kind": "prefill", + "causal": 0, + "hq": hq, + "hkv": hkv, + "d": d, + "window": 0, + "dtype": a.dtype, + "block_size": bs, + "source": jit.source_digest(prefill=True), + "family": gpu_family(), + "cells": cells, + } + meta = { + "vllm": vllm.__version__, + "torch": torch.__version__, + "host": platform.node(), + } + + def variant(knobs): + return R.PrefillVariant( + d, hq, hkv, 0, bs, NON_CAUSAL_STEPS[0], **knobs, dtype=dtype, causal=0 + ) + + def pid(knobs): + return json.dumps(sorted(knobs.items())) + + current = { + m: R.prefill_variant_for(hq, hkv, d, bs, 1, 0, dtype, m, causal=False) + for m in NON_CAUSAL_STEPS + } + raw = ( + [json.loads(x) for x in Path(a.candidates).read_text().splitlines() if x] + if a.candidates + else non_causal_grid(d) + ) + cands: dict[str, R.PrefillVariant] = {} + for k in [*raw, *(prefill_knobs_of(v) for v in current.values() if v)]: + try: + cands.setdefault(pid(k), variant(k)) + except (TypeError, ValueError): + continue + items = list(cands.items())[: a.limit or None] + known: dict[str, dict] = {} + if a.reuse and Path(a.points).exists(): + for line in Path(a.points).read_text().splitlines(): + try: + rec = json.loads(line) + except json.JSONDecodeError: + continue + if all(rec.get(k) == v for k, v in fixed.items()): + known[pid(rec["knobs"])] = rec + t0 = time.time() + # Reused points run again as finalists. + jit.precompile([v for _, v in items]) + jit.seal() + print( + f"# {tag}: {len(items)} points, built in {time.time() - t0:.0f} s", flush=True + ) + time.sleep(5) # let clocks settle after the build + + dev = torch.device("cuda") + data = {} + + def inputs(m): + if m not in data: + torch.manual_seed(m) + q = (torch.randn(m, hq, d, device=dev) * 0.5).to(dtype) + k = (torch.randn(m, hkv, d, device=dev) * 0.5).to(dtype) + v = (torch.randn(m, hkv, d, device=dev) * 0.5).to(dtype) + sl = torch.tensor([m], device=dev, dtype=torch.int32) + ref = torch.empty_like(q) + data[m] = (q, k, v, sl, ref) + return data[m] + + def triton(m): + q, k, v, _, ref = inputs(m) + starts = torch.zeros(1, device=dev, dtype=torch.int32) + lens = torch.tensor([m], device=dev, dtype=torch.int32) + return lambda: context_attention_fwd( + q, k, v, ref, b_start_loc=starts, b_seq_len=lens, max_input_len=m, + is_causal=False, softmax_scale=d**-0.5, + ) # fmt: skip + + scratch = {} + + def op(module, var, m, out): + q, k, v, sl, _ = inputs(m) + key = var.scratch_rows(dev) + if key not in scratch: + scratch[key] = R.make_prefill_scratch([var], dev) + return lambda: module.encoder_attn(q, k, v, out, *scratch[key], sl, m, d**-0.5) + + for m in cells: + triton(m)() + torch.accelerator.synchronize() + + # 1. Check and time every point. + with open(a.points, "a") as fh: + for i, (key, v) in enumerate(items): + if key in known: + continue + rec = { + **fixed, + "knobs": prefill_knobs_of(v), + **meta, + "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + module = jit.load(v) + if module is None: + rec.update(status="build", error=(jit.failure(v) or "")[:400]) + else: + us = [] + for m in cells: + out = torch.full_like(inputs(m)[0], float("nan")) + run = op(module, v, m, out) + run() + torch.accelerator.synchronize() + err = (out.float() - inputs(m)[4].float()).abs().max() + if not err.item() <= tol: + rec.update(status="wrong", cell=m, err=err.item()) + break + us.append( + do_bench(run, warmup=5, rep=40, return_mode="median") * 1e3 + ) + else: + rec.update(status="ok", us=[round(u, 2) for u in us]) + fh.write(json.dumps(rec) + "\n") + fh.flush() + known[key] = rec + if (i + 1) % 20 == 0: + print( + f"# {i + 1}/{len(items)} points, " + f"{(time.time() - t0) / 60:.1f} min", + flush=True, + ) + ok = { + k: dict(zip(cells, r["us"])) + for k, r in known.items() + if r.get("status") == "ok" + } + if not ok: + print(f"{tag}: no point measured correct", flush=True) + return + + def bucket(m): + return max(x for x in NON_CAUSAL_STEPS if x <= m) + + step_cells = {s_: [c for c in cells if bucket(c) == s_] for s_ in NON_CAUSAL_STEPS} + + # 2. Finalists of each step, re-timed against Triton in turn. + finalists: dict[int, set[str]] = {} + for st in NON_CAUSAL_STEPS: + sc = step_cells[st] + fastest = {c: min(t[c] for t in ok.values()) for c in sc} + ranked = sorted( + ok, + key=lambda k: -weighted_speedup( + [weight[c] for c in sc], [fastest[c] / ok[k][c] for c in sc] + ), + ) + finalists[st] = set(ranked[: a.top]) + if current[st] is not None and pid(prefill_knobs_of(current[st])) in ok: + finalists[st].add(pid(prefill_knobs_of(current[st]))) + for st, prev in zip(NON_CAUSAL_STEPS[1:], NON_CAUSAL_STEPS): + finalists[st] |= finalists[prev] + times: dict = {} + for _ in range(a.confirm_rounds): + for st in NON_CAUSAL_STEPS: + for c in step_cells[st]: + times.setdefault(("triton", c), []).append( + do_bench(triton(c), warmup=5, rep=40, return_mode="median") * 1e3 + ) + for k in sorted(finalists[st]): + out = torch.empty_like(inputs(c)[0]) + run = op(jit.load(cands[k]), cands[k], c, out) + times.setdefault((k, c), []).append( + do_bench(run, warmup=5, rep=40, return_mode="median") * 1e3 + ) + med = {key: statistics.median(t) for key, t in times.items()} + + # 3. Per step, prefer points never slower than Triton, then speedup. + steps, prev = [], None + for st in NON_CAUSAL_STEPS: + sc = step_cells[st] + w = [weight[c] for c in sc] + scored = {} + for k in finalists[st]: + ratio = [med[("triton", c)] / med[(k, c)] for c in sc] + scored[k] = (min(ratio) >= 1.0, weighted_speedup(w, ratio), ratio) + best = max(scored, key=lambda k: scored[k][:2]) + if ( + prev in scored + and (scored[prev][0] or not scored[best][0]) + and scored[prev][1] * PREFILL_SLACK >= scored[best][1] + ): + best = prev + cur = None if current[st] is None else prefill_knobs_of(current[st]) + steps.append( + { + "m": st, + "current": cur, + "winner": dict(json.loads(best)), + "speedup": scored[best][1], + "worst": min(scored[best][2]), + "cells": { + str(c): { + "ours": med[(best, c)], + "base": med[("triton", c)], + "roof": roof[c] / med[(best, c)], + } + for c in sc + }, + } + ) + prev = best + rf = [roof[c] / med[(best, c)] for c in sc] + print( + f"{tag} M>={st}: {prefill_label(steps[-1]['winner'])} " + f"speedup {steps[-1]['speedup']:.2f}x, worst cell " + f"{steps[-1]['worst']:.2f}x, %roof {min(rf) * 100:.1f}-{max(rf) * 100:.1f}", + flush=True, + ) + decision = { + **fixed, + "current": {st: s_["current"] for st, s_ in zip(NON_CAUSAL_STEPS, steps)}, + "steps": steps, + "accepted": True, + **meta, + "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + with open(a.decisions, "a") as out_: + out_.write(json.dumps(decision) + "\n") + + def apply_prefill(dec: dict, block_sizes: list[int]) -> int: """Write one prefill decision into the prefill configs.""" import jit @@ -1134,13 +1456,14 @@ def apply_prefill(dec: dict, block_sizes: list[int]) -> int: if dec["source"] != jit.source_digest(prefill=True): print(f"{tag}: refused, measured on another kernel source") return 0 + causal = dec.get("causal", 1) family = dec.get("family") or gpu_family() - lines, rows = read_configs(configs_h("prefill"), family) + lines, rows = read_configs(configs_h("prefill"), family, R.PREFILL_CONFIG_FIELDS) col = {f: i for i, f in enumerate(R.PREFILL_CONFIG_FIELDS)} dtype = utils.torch_dtype(dec["dtype"]) for step in dec["steps"]: now = R.prefill_variant_for( - hq, hkv, d, dec["block_size"], 1, win, dtype, step["m"] + hq, hkv, d, dec["block_size"], 1, win, dtype, step["m"], causal ) if step["current"] != (None if now is None else prefill_knobs_of(now)): print(f"{tag}: refused, the rows changed since they were measured") @@ -1152,14 +1475,17 @@ def mine(r): r[col["num_q_heads"]], r[col["num_kv_heads"]], r[col["window"]], - ) == (d, hq, hkv, win) + r[col["causal"]], + ) == (d, hq, hkv, win, causal) sizes = sorted({r[col["page_size"]] for r in rows if mine(r)}) or block_sizes kept = [r for r in rows if not mine(r)] # A decided step's winner serves M from its step up to the next step of - # STEPS; outside them the entries stay. + # its grid (NON_CAUSAL_STEPS for a non-causal build); outside them the + # entries stay. decided = {s["m"]: s["winner"] for s in dec["steps"]} - ends = dict(zip(STEPS, (*STEPS[1:], None))) + grid = STEPS if causal else NON_CAUSAL_STEPS + ends = dict(zip(grid, (*grid[1:], None))) written = 0 for bf16, bs in itertools.product((0, 1), sizes): group = sorted( @@ -1194,6 +1520,7 @@ def mine(r): b, **knobs, dtype=torch.bfloat16 if bf16 else torch.float16, + causal=causal, ) kept.append(R.prefill_config_fields(v)) written += 1 @@ -1217,13 +1544,36 @@ def mine(r): BATCH_WORKING_SET = 96 << 20 -def prefill_batch_cells(window, kv_bytes_per_key): - """((min mean M, min keys), representative batches) of every cell.""" +def prefill_batch_cells(window, kv_bytes_per_key, causal=True): + """((min mean M, min keys), representative batches) of every cell. + + Non-causal (an encoder's segments, L = S): a batch's mean M and keys are + both L, so only the cells holding some L >= PREFILL_MIN_M are reached; + each takes batches of a total near BATCH_TOKENS and a quarter of it. + """ import numpy as np import vllm.v1.attention.ops.rocm_rdna_hip as R out = [] + if not causal: + ms = (*R.PREFILL_BATCH_MEAN_M, 1 << 30) + ks = (*R.PREFILL_BATCH_KEYS, 1 << 30) + for i, m_lo in enumerate(R.PREFILL_BATCH_MEAN_M): + for j, k_lo in enumerate(R.PREFILL_BATCH_KEYS): + lo = max(m_lo, k_lo, R.PREFILL_MIN_M) + hi = min(ms[i + 1], ks[j + 1], 16384) + if lo >= hi: + continue + length = min(hi - 1, max(lo, round(math.sqrt(lo * hi)))) + n = max(2, BATCH_TOKENS // length) + specs = [] + for nseq in dict.fromkeys((n, max(2, n // 4))): + starts = np.arange(nseq + 1) * length + lens = np.full(nseq, length) + specs.append((starts, lens, f"{nseq}q{length}s{length}")) + out.append(((m_lo, k_lo), specs)) + return out for i, m_lo in enumerate(R.PREFILL_BATCH_MEAN_M): m = BATCH_M[m_lo] for j, k_lo in enumerate(R.PREFILL_BATCH_KEYS): @@ -1253,9 +1603,15 @@ def tune_prefill_batch(a) -> None: jit.install() dtype = utils.torch_dtype(a.dtype) hq, hkv, d, win, bs = a.hq, a.hkv, a.head_dim, a.window, a.block_size - tag = f"{hq}/{hkv}/{d}" + (f" w{win}" if win else "") + f" batch {a.dtype}" + causal = not a.non_causal + tag = ( + f"{hq}/{hkv}/{d}" + + (f" w{win}" if win else "") + + ("" if causal else " non-causal") + + f" batch {a.dtype}" + ) dev = torch.device("cuda") - rows = R.prefill_variants(hq, hkv, d, bs, win, dtype) + rows = R.prefill_variants(hq, hkv, d, bs, win, dtype, causal) if not rows: raise SystemExit(f"{tag}: no prefill rows") # Rows differing in max_segments or head_group build the same kernel. @@ -1288,12 +1644,18 @@ def timed(launches, iters=25): cells = [] kv_bytes_per_key = hkv * 2 * d * torch.finfo(dtype).bits // 8 - for (m_lo, k_lo), specs in prefill_batch_cells(win, kv_bytes_per_key): + for (m_lo, k_lo), specs in prefill_batch_cells(win, kv_bytes_per_key, causal): times: dict = {} for starts, lens, spec in specs: torch.manual_seed(0) nseq, pages = len(lens), -(-int(lens[0]) // bs) - kv = (torch.randn(nseq * pages, hkv, bs, 2 * d, device=dev) * 0.5).to(dtype) + # Causal: a paged cache; non-causal: K and V (in place). + shape = ( + (nseq * pages, hkv, bs, 2 * d) + if causal + else (2, int(starts[-1]), hkv, d) + ) + kv = (torch.randn(shape, device=dev) * 0.5).to(dtype) kvs = [kv] + [ kv.clone() for _ in range(BATCH_WORKING_SET // (kv.numel() * kv.element_size())) @@ -1325,6 +1687,21 @@ def timed(launches, iters=25): query_start_loc=qsl, head_group=hg, ) + if causal + else functools.partial( + mod.encoder_attn, + q, + cache[0], + cache[1], + out, + *scratch, + sl, + int(lens.max()), + d**-0.5, + plan=plan, + query_start_loc=qsl, + head_group=hg, + ) for cache in kvs ] ) @@ -1359,6 +1736,7 @@ def timed(launches, iters=25): "hkv": hkv, "d": d, "window": win, + "causal": int(causal), "dtype": a.dtype, "block_size": bs, "source": jit.source_digest(prefill=True), @@ -1384,21 +1762,30 @@ def apply_prefill_batch(dec: dict, block_sizes: list[int]) -> int: import vllm.v1.attention.ops.rocm_rdna_hip as R hq, hkv, d, win = dec["hq"], dec["hkv"], dec["d"], dec["window"] + causal = dec.get("causal", 1) tag = f"{hq}/{hkv}/{d}" + (f" w{win}" if win else "") + " batch" - rows = R.prefill_variants(hq, hkv, d, dec["block_size"], win, torch.float16) + rows = R.prefill_variants(hq, hkv, d, dec["block_size"], win, torch.float16, causal) now = [prefill_knobs_of(v) | {"min_query_len": v.min_query_len} for v in rows] if now != dec["rows"]: print(f"{tag}: refused, the prefill builds changed since it was measured") return 0 family = dec.get("family") or gpu_family() - lines, table = read_configs(configs_h("prefill_batch"), family) + lines, table = read_configs( + configs_h("prefill_batch"), family, R.PREFILL_BATCH_CONFIG_FIELDS + ) pages = ( sorted( - {k[3] for k in R._prefill_table() if k[:3] == (hq, hkv, d) and k[4] == win} + { + k[3] + for k in R._prefill_table() + if k[:3] == (hq, hkv, d) and k[4] == win and k[6] == causal + } ) or block_sizes ) - table = [r for r in table if (r[0], r[1], r[2], r[3]) != (d, hq, hkv, win)] + table = [ + r for r in table if (r[0], r[1], r[2], r[3], r[9]) != (d, hq, hkv, win, causal) + ] for page in pages: for c in dec["cells"]: table.append( @@ -1412,6 +1799,7 @@ def apply_prefill_batch(dec: dict, block_sizes: list[int]) -> int: c["min_keys"], c["row"], c["head_group"], + causal, ] ) write_configs( @@ -1470,16 +1858,33 @@ def _table_lines(lines: list[str], family: str) -> range: return range(start, end) -def read_configs(path: Path, family: str) -> tuple[list[str], list[list[int]]]: - """The lines of a configs header, and the entries of family's table.""" +def read_configs( + path: Path, family: str, names: tuple[str, ...] | None = None +) -> tuple[list[str], list[list[int]]]: + """The lines of a configs header, and the entries of family's table; with + `names`, left out trailing fields read as their defaults.""" import jit lines = path.read_text().splitlines() at = _table_lines(lines, family) rows = [jit.parse_cpp_fields(m[1]) for i in at if (m := _ENTRY.match(lines[i]))] + if names: + tail = _trailing_defaults(names) + rows = [r + tail[len(tail) - (len(names) - len(r)) :] for r in rows] return lines, rows +def _trailing_defaults(names: tuple[str, ...]) -> list[int]: + """Defaults of the fields an entry may leave out, at the end of `names`.""" + defaults = {"causal": 1, "row_tiles": 1} + tail: list[int] = [] + for name in reversed(names): + if name not in defaults: + break + tail.insert(0, defaults[name]) + return tail + + def write_configs( path: Path, lines: list[str], @@ -1492,7 +1897,20 @@ def write_configs( import jit at = _table_lines(lines, family) - entries = [f" {{{jit.cpp_fields(r, names)}}}," for r in sorted(rows)] + # Trailing fields at their defaults are left out. + tail = _trailing_defaults(names) + first_tail = len(names) - len(tail) + + def kept(r): + n = len(r) + while n > first_tail and r[n - 1] == tail[n - 1 - first_tail]: + n -= 1 + return n + + entries = [ + f" {{{jit.cpp_fields(r[: kept(r)], names[: kept(r)])}}}," + for r in sorted(rows) + ] text = lines[: at.start] + entries + lines[at.stop :] path.write_text("\n".join(text) + "\n") @@ -1610,6 +2028,12 @@ def main() -> None: help="tune the prefill batch configs: the prefill build and " "head group of every cell of one configuration's batches", ) + p.add_argument( + "--non-causal", + action="store_true", + help="with --prefill or --prefill-batch: every query token sees every " + "key of its sequence (an encoder's segments, L = S)", + ) p.add_argument( "--prefixes", type=int, @@ -1679,6 +2103,8 @@ def main() -> None: p.error("tuning needs --" + ", --".join(n.replace("_", "-") for n in missing)) if a.prefill_batch: tune_prefill_batch(a) + elif a.prefill and a.non_causal: + tune_prefill_non_causal(a) elif a.prefill: tune_prefill(a) else: diff --git a/benchmarks/kernels/rocm_rdna_hip_attn/tools/utils.py b/benchmarks/kernels/rocm_rdna_hip_attn/tools/utils.py index c277009f4474..2f19bdb6a8d3 100644 --- a/benchmarks/kernels/rocm_rdna_hip_attn/tools/utils.py +++ b/benchmarks/kernels/rocm_rdna_hip_attn/tools/utils.py @@ -68,6 +68,7 @@ def _arch() -> str: "key_waves", "max_segments", "head_group", + "row_tiles", ) @@ -180,6 +181,30 @@ def roofline_us( return roofline(hq, hkv, d, m, s, itemsize, block_size, window)[0] +def roofline_noncausal( + hq: int, hkv: int, d: int, segs: list[int], itemsize: int = 2 +) -> tuple[float, str]: + """The floor of one non-causal call over sequences of `segs` tokens, each + query seeing every key of its sequence, with the WMMA prefill model: Q, K, + V and the output at peak bandwidth against the 16x16 score tiles at peak + throughput, plus a dispatch and three DRAM round trips. + + Returns: + (microseconds, "memory" or "compute"). + + """ + t, g = WMMA_TILE, hq // hkv + n = sum(segs) + moved = (2 * n * hq + 2 * n * hkv) * d * itemsize + memory_us = moved / (PEAK_GIBS * 1024**3) * 1e6 + tiles = sum(-(-s * g // t) * -(-s // t) for s in segs) * hkv + compute_us = tiles * 4 * d * t**2 / (PEAK_TFLOPS * 1e6) + ( + tiles * SOFTMAX_CYCLES / (SIMDS * WMMA_GHZ * 1e3) + ) + bound = "compute" if compute_us > memory_us else "memory" + return DISPATCH_US + 3 * DRAM_LATENCY_US + max(memory_us, compute_us), bound + + def add_dtype_argument(p: argparse.ArgumentParser) -> None: p.add_argument( "--dtype", diff --git a/cmake/rocm_rdna_attn.cmake b/cmake/rocm_rdna_attn.cmake index cdffcd45db06..2c4f5bd1d876 100644 --- a/cmake/rocm_rdna_attn.cmake +++ b/cmake/rocm_rdna_attn.cmake @@ -10,7 +10,7 @@ # function(rdna_attn_target TARGET FAMILY ARCHES) set(_KERNELS decode decode_batch prefill) - set(_HEAD_DIMS 64 96 128 256 512) + set(_HEAD_DIMS 32 48 64 72 80 96 128 256 512) set(_DTYPES Fp16 Bf16) string(TOLOWER ${FAMILY} _FAMILY_NAME) set(_OUT ${CMAKE_CURRENT_BINARY_DIR}/rdna_attn/${_FAMILY_NAME}) diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index bf5aa4c6e59a..e76538710835 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -164,6 +164,14 @@ void rdna_prefill_attn(int64_t config, torch::Tensor& q, const std::optional& plan, const std::optional& query_start_loc, int64_t head_group); +void rdna_encoder_attn(int64_t config, torch::Tensor& q, torch::Tensor& k, + torch::Tensor& v, torch::Tensor& out, + torch::Tensor& partial_o, torch::Tensor& partial_ml, + torch::Tensor& counters, torch::Tensor& seq_lens, + int64_t max_seq_len, double scale, + const std::optional& plan, + const std::optional& query_start_loc, + int64_t head_group); void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, diff --git a/csrc/rocm/rdna_attn_launch.h b/csrc/rocm/rdna_attn_launch.h index 3a9a2a34ce90..7701b4b8c79b 100644 --- a/csrc/rocm/rdna_attn_launch.h +++ b/csrc/rocm/rdna_attn_launch.h @@ -67,6 +67,11 @@ struct PrefillArgs { int num_entries; int bt_stride; int head_group; + // Non-causal builds: V, and the token strides of Q, K (in `kv`) and V. + const void* v; + int q_stride; + int k_stride; + int v_stride; }; using PrefillFn = void (*)(const PrefillArgs&); diff --git a/csrc/rocm/rdna_attn_types.h b/csrc/rocm/rdna_attn_types.h index 18d957c9518a..5ba3fe60ac34 100644 --- a/csrc/rocm/rdna_attn_types.h +++ b/csrc/rocm/rdna_attn_types.h @@ -5,7 +5,11 @@ namespace rdna { enum class HeadDim : int { + k32 = 32, + k48 = 48, k64 = 64, + k72 = 72, + k80 = 80, k96 = 96, k128 = 128, k256 = 256, diff --git a/csrc/rocm/rdna_decode_attn.cu b/csrc/rocm/rdna_decode_attn.cu index f400fdf2b43e..5a2deafdd5c9 100644 --- a/csrc/rocm/rdna_decode_attn.cu +++ b/csrc/rocm/rdna_decode_attn.cu @@ -40,6 +40,23 @@ struct u3w { __device__ unsigned operator[](int i) const { return w[i]; } }; +// Two dims (D=32): one dword. +struct u1w { + unsigned w[1]; + __device__ unsigned& operator[](int i) { return w[i]; } + __device__ unsigned operator[](int i) const { return w[i]; } +}; + +// Dims per lane of the dot path: the existing rule where D is a multiple of +// 32, else the fewest whole dot2 pairs whose lanes per key divide the wave. +constexpr int dot_dims_per_lane(int headDim) { + if (headDim == 128) return 8; + if (headDim % 32 == 0) return std::max(2, headDim / 32); + for (int dims = 2; dims <= headDim; dims += 2) + if (headDim % dims == 0 && 32 % (headDim / dims) == 0) return dims; + return headDim; +} + template struct Decode { static constexpr bool kBatched = kBatchedBuild; @@ -137,8 +154,7 @@ struct Decode { // evaluated for every build; the asserts only bind its own. // One workgroup per (q head, segment); each wave takes kTileKeys keys at a // time with the head dim across its lanes. - static constexpr int kDimsPerLane = - kHeadDim == 128 ? 8 : kHeadDim / kWaveSize; + static constexpr int kDimsPerLane = dot_dims_per_lane(kHeadDim); static constexpr int kLanesPerKey = kHeadDim / kDimsPerLane; static constexpr int kKeysPerStep = kWaveSize / kLanesPerKey; static constexpr int kKeysPerLane = std::max(4 / kKeysPerStep, 1); @@ -474,10 +490,19 @@ struct Decode { static constexpr int kPaddedRows = kRowTiles * 16; static constexpr int kKeyTiles = kWaves / kHeadDimSplit; static constexpr int kBlockKeys = 16 * kKeyTiles; - static constexpr int kPartDims = kHeadDim / kHeadDimSplit; + // The head dim padded to whole 16-dim chunks (D=72: 80); Q and the K tile + // hold zeros past D, the odd lane slice only real dims. + static constexpr int kHeadDimP = (kHeadDim + 15) / 16 * 16; + static constexpr int kPartDims = kHeadDimP / kHeadDimSplit; static constexpr int kPartChunks = kPartDims / 16; - // Consecutive dims of a K or V row per lane. + // Dims of a K or V row per lane: kLaneA consecutive ones (dims + // kLaneA * r .. of row tile lane r) and, when kLaneDims is odd, one more + // (dim 16 * kLaneA + r), so D=48 runs as 32 + 16 and D=80 as 64 + 16. static constexpr int kLaneDims = kPartDims / 16; + static constexpr int kLaneB = kLaneDims % 2; + static constexpr int kLaneA = kLaneDims - kLaneB; + // Real dims of the odd slice: 16, or 8 for D=72. + static constexpr int kTailDims = kLaneB ? kHeadDim - 16 * kLaneA : 0; // P goes into P@V as a high and a low part for precision. With at most 8 rows // the low part uses the padding columns of the same WMMA; past that it costs // a second WMMA, which gfx12 skips for fp16. @@ -488,7 +513,7 @@ struct Decode { static constexpr int kTilesPerWave = kVInLds ? 2 : 1; static constexpr int kWaveTileBytes = kTilesPerWave * 16 * kTileRowBytes; static constexpr int kQRowPad = 8; - static constexpr int kQBytes = kPaddedRows * (kHeadDim + kQRowPad) * 2; + static constexpr int kQBytes = kPaddedRows * (kHeadDimP + kQRowPad) * 2; // Merge buffer: one slot of real rows per (live tile, part). static constexpr int kMergeRowStride = kPartDims + 4; static constexpr int kMergeSlotFloats = @@ -523,11 +548,13 @@ struct Decode { "row groups split the rows evenly"); static_assert(kDotProduct || kWaves % kHeadDimSplit == 0, "whole tiles per workgroup"); - static_assert(kDotProduct || kHeadDim / 16 % kHeadDimSplit == 0, + static_assert(kDotProduct || kHeadDimP / 16 % kHeadDimSplit == 0, "whole chunks per part"); - static_assert(kDotProduct || kLaneDims == 4 || kLaneDims == 6 || - kLaneDims == 8, - "a lane's K and V slice is one b64, b96 or b128"); + static_assert(kDotProduct || kLaneA == 2 || kLaneA == 4 || kLaneA == 6 || + kLaneA == 8, + "a lane's K and V slice is one b32, b64, b96 or b128"); + static_assert(kDotProduct || !kLaneB || kHeadDimSplit == 1, + "an odd lane slice is not split over the head dim"); static_assert(kDotProduct || kKeyTiles <= kWaveSize, "one lane per tile loads its page"); static_assert(kDotProduct || kHeadDimSplit == 1 || @@ -536,11 +563,12 @@ struct Decode { static_assert(kDotProduct || !kVInLds || !kPrefetch, "V_IN_LDS replaces PREFETCH's second tile"); - using vrow_t = - std::conditional_t >; + using vrow_t = std::conditional_t< + kLaneA == 8, u4v, + std::conditional_t > >; // A six-dim slice is only 8-byte aligned. - static constexpr int kRowChunk = kLaneDims % 4 ? 2 : 4; + static constexpr int kRowChunk = kLaneA % 4 ? 2 : 4; typedef float fchunk __attribute__((ext_vector_type(kRowChunk))); // A wave's share of a 16-key tile: half h holds keys 2e + h, each lane @@ -548,6 +576,9 @@ struct Decode { struct Tile { vrow_t k[8]; vrow_t v[8]; + // The odd dim, when kLaneB. + unsigned short kt[8]; + unsigned short vt[8]; }; // Lane t holds the page of tile t of a block, clamped to the last page. @@ -584,6 +615,15 @@ struct Decode { [](unsigned x) { return xhalf_u(x); }); } + // A = V^T for the odd dim (kLaneB). + static __device__ __forceinline__ op_t vt_frag(const unsigned short* v) { + return Elem::own_first_op( + [&](int i) { + return (unsigned)v[2 * i] | ((unsigned)v[2 * i + 1] << 16); + }, + [](unsigned x) { return xhalf_u(x); }); + } + // Store a wave's real rows of one row tile, `stride` floats apart. static __device__ __forceinline__ void store_rows( float* dst, const f8* accum, int laneRow, int laneHalf, @@ -591,13 +631,16 @@ struct Decode { float* row = dst + laneRow * stride; #pragma unroll for (int e = 0; e < 8; ++e) { - float values[kLaneDims]; + float values[kLaneA]; #pragma unroll - for (int j = 0; j < kLaneDims; ++j) values[j] = accum[j][e]; + for (int j = 0; j < kLaneA; ++j) values[j] = accum[j][e]; #pragma unroll - for (int k = 0; k < kLaneDims; k += kRowChunk) - *(fchunk*)(row + kLaneDims * acc_row(e, laneHalf) + k) = + for (int k = 0; k < kLaneA; k += kRowChunk) + *(fchunk*)(row + kLaneA * acc_row(e, laneHalf) + k) = *(const fchunk*)(values + k); + if constexpr (kLaneB) + if (kTailDims == 16 || acc_row(e, laneHalf) < kTailDims) + row[16 * kLaneA + acc_row(e, laneHalf)] = accum[kLaneA][e]; } } @@ -605,14 +648,18 @@ struct Decode { int laneRow, int laneHalf) { const float* row = src + laneRow * kMergeRowStride; #pragma unroll - for (int e = 0; e < 8; ++e) + for (int e = 0; e < 8; ++e) { #pragma unroll - for (int k = 0; k < kLaneDims; k += kRowChunk) { + for (int k = 0; k < kLaneA; k += kRowChunk) { const fchunk values = - *(const fchunk*)(row + kLaneDims * acc_row(e, laneHalf) + k); + *(const fchunk*)(row + kLaneA * acc_row(e, laneHalf) + k); #pragma unroll for (int j = 0; j < kRowChunk; ++j) accum[k + j][e] += values[j]; } + if constexpr (kLaneB) + if (kTailDims == 16 || acc_row(e, laneHalf) < kTailDims) + accum[kLaneA][e] += row[16 * kLaneA + acc_row(e, laneHalf)]; + } } // The first block's pages, read before S is known. @@ -697,7 +744,7 @@ struct Decode { : 0; // The merge reuses the Q and K tile LDS. - auto qShared = reinterpret_cast(lds); + auto qShared = reinterpret_cast(lds); float* mergeShared = reinterpret_cast(lds); char* waveTile = lds + kQBytes + wave * kTilesPerWave * 16 * kTileRowBytes; // Partial scores reuse the wave's K tile. @@ -710,7 +757,11 @@ struct Decode { float* const groupSum = groupMax + kPaddedRows; const elem_t* kvHeadBase = kv + (size_t)kvHead * (kPageSize * kKvRowElems); - const int laneDim = dimPart * kPartDims + kLaneDims * laneRow; + const int laneDim = dimPart * kPartDims + kLaneA * laneRow; + // Lanes past the real tail dims load a clamped one and zero it. + const int laneTailDim = + dimPart * kPartDims + 16 * kLaneA + + (kTailDims == 16 || !kLaneB ? laneRow : laneRow % kTailDims); const int firstChunk = dimPart * kPartChunks; // Load this wave's tile of `block` and the pages of the next block. @@ -734,6 +785,17 @@ struct Decode { for (int e = 0; e < 8; ++e) tile.k[e] = (kAblate & 4) ? vrow_t{} : *(const vrow_t*)(src + 2 * e * kKvRowElems); + if constexpr (kLaneB) { + const elem_t* tail = src - laneDim + laneTailDim; +#pragma unroll + for (int e = 0; e < 8; ++e) { + tile.vt[e] = + *(const unsigned short*)(tail + 2 * e * kKvRowElems + kHeadDim); + tile.kt[e] = *(const unsigned short*)(tail + 2 * e * kKvRowElems); + if (kTailDims < 16 && laneRow >= kTailDims) + tile.vt[e] = tile.kt[e] = 0; + } + } // Keep the loads from sinking to their first use. asm volatile("" ::: "memory"); }; @@ -747,6 +809,13 @@ struct Decode { if (elem < kPaddedRows * kHeadDim) *(e8*)&qShared[elem / kHeadDim][elem % kHeadDim] = qLoads[it]; } + if constexpr (kHeadDimP != kHeadDim) { + // Zero Q past D so the padded chunk adds nothing. + for (int row = thread; row < kPaddedRows; row += kThreads) +#pragma unroll + for (int d = kHeadDim; d < kHeadDimP; d += 8) + *(e8*)&qShared[row][d] = e8{}; + } const float scaleLog2 = scale * kLog2e; const int context = seqLen - kMaxQueryLen; @@ -778,22 +847,34 @@ struct Decode { #pragma unroll for (int e = 0; e < 8; ++e) { const int offset = - (2 * e + laneHalf) * kTileRowBytes + kLaneDims * laneRow * 2; + (2 * e + laneHalf) * kTileRowBytes + kLaneA * laneRow * 2; *(vrow_t*)(waveTile + offset) = tile.k[e]; *(vrow_t*)(waveTile + 16 * kTileRowBytes + offset) = tile.v[e]; + if constexpr (kLaneB) { + const int tailOffset = + (2 * e + laneHalf) * kTileRowBytes + (16 * kLaneA + laneRow) * 2; + *(unsigned short*)(waveTile + tailOffset) = tile.kt[e]; + *(unsigned short*)(waveTile + 16 * kTileRowBytes + tailOffset) = + tile.vt[e]; + } } asm volatile("" ::: "memory"); }; auto stageK = [&](const Tile& tile) { #pragma unroll - for (int e = 0; e < 8; ++e) + for (int e = 0; e < 8; ++e) { *(vrow_t*)(waveTile + (2 * e + laneHalf) * kTileRowBytes + - kLaneDims * laneRow * 2) = tile.k[e]; + kLaneA * laneRow * 2) = tile.k[e]; + if constexpr (kLaneB) + *(unsigned short*)(waveTile + (2 * e + laneHalf) * kTileRowBytes + + (16 * kLaneA + laneRow) * 2) = tile.kt[e]; + } asm volatile("" ::: "memory"); }; // Q@K, softmax update and P@V for the tile staged at kTile. - auto process = [&](Tile& tile, vrow_t* v, const char* kTile, int block) { + auto process = [&](Tile& tile, vrow_t* v, unsigned short* vt, + const char* kTile, int block) { const int tileKey = block * kBlockKeys + keyTile * 16; if constexpr (kAblate & 64) { unsigned x = 0; @@ -845,13 +926,21 @@ struct Decode { lds_barrier(); } vrow_t ldsV[8]; + unsigned short ldsVt[8]; if constexpr (kVInLds) { #pragma unroll - for (int e = 0; e < 8; ++e) + for (int e = 0; e < 8; ++e) { ldsV[e] = *(const vrow_t*)(kTile + 16 * kTileRowBytes + (2 * e + laneHalf) * kTileRowBytes + - kLaneDims * laneRow * 2); + kLaneA * laneRow * 2); + if constexpr (kLaneB) + ldsVt[e] = + *(const unsigned short*)(kTile + 16 * kTileRowBytes + + (2 * e + laneHalf) * kTileRowBytes + + (16 * kLaneA + laneRow) * 2); + } v = ldsV; + vt = ldsVt; } // Zero V past S and before the window (possibly freed pages), so no @@ -859,14 +948,19 @@ struct Decode { if (__builtin_expect(tileKey + 16 > seqLen, 0)) { #pragma unroll for (int e = 0; e < 8; ++e) - if (tileKey + 2 * e + laneHalf >= seqLen) v[e] = vrow_t{}; + if (tileKey + 2 * e + laneHalf >= seqLen) { + v[e] = vrow_t{}; + if constexpr (kLaneB) vt[e] = 0; + } } if constexpr (kWindow) { if (__builtin_expect(tileKey < context - (kWindow - 1), 0)) { #pragma unroll for (int e = 0; e < 8; ++e) - if (tileKey + 2 * e + laneHalf < context - (kWindow - 1)) + if (tileKey + 2 * e + laneHalf < context - (kWindow - 1)) { v[e] = vrow_t{}; + if constexpr (kLaneB) vt[e] = 0; + } } } @@ -939,7 +1033,8 @@ struct Decode { for (int j = 0; j < kLaneDims; ++j) { accum[rt][j] *= accumScale; if ((kAblate & 2) && j) continue; - const op_t vOperand = v_frag(v, j); + const op_t vOperand = + (kLaneB && j == kLaneA) ? vt_frag(vt) : v_frag(v, j); accum[rt][j] = wmma(vOperand, pOperand, accum[rt][j]); if constexpr (kLowWmma) accum[rt][j] = wmma(vOperand, pLowOperand, accum[rt][j]); @@ -954,25 +1049,25 @@ struct Decode { block += 2 * numSegments) { stageK(tileA); if (block + numSegments < numBlocks) issue(tileB, block + numSegments); - process(tileA, tileA.v, waveTile, block); + process(tileA, tileA.v, tileA.vt, waveTile, block); if (block + numSegments >= numBlocks) break; stageK(tileB); if (block + 2 * numSegments < numBlocks) issue(tileA, block + 2 * numSegments); - process(tileB, tileB.v, waveTile, block + numSegments); + process(tileB, tileB.v, tileB.vt, waveTile, block + numSegments); } } else if constexpr (kVInLds) { for (int block = firstBlock + segment; block < numBlocks; block += numSegments) { stageKv(tileA); if (block + numSegments < numBlocks) issue(tileA, block + numSegments); - process(tileA, nullptr, waveTile, block); + process(tileA, nullptr, nullptr, waveTile, block); } } else { for (int block = firstBlock + segment; block < numBlocks; block += numSegments) { stageK(tileA); - process(tileA, tileA.v, waveTile, block); + process(tileA, tileA.v, tileA.vt, waveTile, block); if (block + numSegments < numBlocks) issue(tileA, block + numSegments); } } @@ -1056,8 +1151,11 @@ struct Decode { #pragma unroll for (int j = 0; j < kLaneDims; ++j) x[j] = to_elem(accum[rt][j][e] * inv); - *(vrow_t*)(dst + kLaneDims * acc_row(e, laneHalf)) = + *(vrow_t*)(dst + kLaneA * acc_row(e, laneHalf)) = *(const vrow_t*)x; + if constexpr (kLaneB) + if (kTailDims == 16 || acc_row(e, laneHalf) < kTailDims) + dst[16 * kLaneA + acc_row(e, laneHalf)] = x[kLaneA]; } } else { store_rows(partialAcc + (partialBase + rt * 16) * kHeadDim + diff --git a/csrc/rocm/rdna_prefill_attn.cu b/csrc/rocm/rdna_prefill_attn.cu index 7b9269fd8e8b..cb5637a3411c 100644 --- a/csrc/rocm/rdna_prefill_attn.cu +++ b/csrc/rocm/rdna_prefill_attn.cu @@ -56,6 +56,8 @@ struct Prefill { static constexpr int kNumQHeads = C.num_q_heads; static constexpr int kNumKvHeads = C.num_kv_heads; static constexpr int kWindow = C.window; + // Non-causal: every query token sees every key of its sequence. + static constexpr bool kCausal = C.causal; static constexpr bool kBf16 = C.dtype == Dtype::kBf16; static constexpr int kPageSize = C.page_size; static constexpr int kWaves = C.waves; @@ -68,6 +70,8 @@ struct Prefill { static constexpr bool kDoubleBuffer = C.double_buffer; static constexpr int kKeyWaves = C.key_waves; static constexpr int kMaxSegments = C.max_segments; + // 16-row tiles per wave: two share every K and V operand read (non-causal). + static constexpr int kRowTiles = C.row_tiles; // Fewest steps per KV segment. static constexpr int kMinSegmentSteps = 16; @@ -76,10 +80,13 @@ struct Prefill { static constexpr int kThreads = kWaves * kWaveSize; static constexpr int kG = kNumQHeads / kNumKvHeads; static constexpr int kRowWaves = kWaves / (kDimParts * kKeyWaves); - static constexpr int kBlockRows = kRowWaves * 16; + static constexpr int kBlockRows = kRowWaves * 16 * kRowTiles; static constexpr int kSubTiles = kKeyTile / 16; static constexpr int kWaveSubTiles = kSubTiles / kKeyWaves; - static constexpr int kPartDims = kHeadDim / kDimParts; + // The head dim padded to whole 16-dim chunks (D=72: 80). Q, the K tile and + // V_T hold zeros past D; loads past D are clamped and stored as zeros. + static constexpr int kHeadDimP = (kHeadDim + 15) / 16 * 16; + static constexpr int kPartDims = kHeadDimP / kDimParts; static constexpr int kPartChunks = kPartDims / 16; // Consecutive dims of a V row per lane. static constexpr int kLaneDims = kPartDims / 16; @@ -87,17 +94,22 @@ struct Prefill { static constexpr int kKvRowElems = 2 * kHeadDim; static constexpr int kPageElems = kPageSize * kNumKvHeads * kKvRowElems; // LDS rows are padded by 16 bytes to avoid bank conflicts. - static constexpr int kRowPitch = kHeadDim + 8; + static constexpr int kRowPitch = kHeadDimP + 8; // V_T row pitch, chosen to avoid bank conflicts: an odd number of 16-byte // units on gfx12, where a 6-unit pitch halves the LDS read rate. static constexpr int kVtPitch = kGfx12 ? kKeyTile + 8 : kKeyTile == 16 ? 16 : kKeyTile + 16; // The V tile: V rows or V^T rows; a placeholder when v_t 2 skips it. - static constexpr int kVTileRows = kVTr ? 1 : kVT ? kHeadDim : kKeyTile; + static constexpr int kVTileRows = kVTr ? 1 : kVT ? kHeadDimP : kKeyTile; static constexpr int kVTilePitch = kVTr ? 8 : kVT ? kVtPitch : kRowPitch; static constexpr int kBuffers = kDoubleBuffer ? 2 : 1; - static constexpr int kRowChunks = kHeadDim / 8; + // Elements per load unit: 8, or 4 where a 16-key sub-tile of 8-element + // units would not be whole waves (D=48, 80). Non-causal builds load 8 of + // the D real dims of K and V in place (the padding is zeroed once). + static constexpr int kUnitElems = !kCausal || kHeadDimP / 8 % 4 == 0 ? 8 : 4; + static constexpr int kRowChunks = + (kCausal ? kHeadDimP : kHeadDim) / kUnitElems; static constexpr int kKUnits = kKeyTile * kRowChunks; // With V_T a V unit is keys kk and kk + 2 of a sub-tile. static constexpr int kVUnits = (kVT ? kKeyTile / 2 : kKeyTile) * kRowChunks; @@ -111,8 +123,10 @@ struct Prefill { kKeyTile > kPageSize ? kKeyTile / kPageSize : 1; static constexpr float kLog2e = 1.44269504088896340736f; static constexpr float kMaxSlack = 8.f; - // D=512 uses buffer loads to save the registers of 64-bit addresses. - static constexpr bool kBufferLoads = kHeadDim >= 512; + // Buffer loads save the registers of 64-bit addresses: D=512, and the + // non-causal builds. + static constexpr bool kBufferLoads = + !kCausal || (kHeadDim >= 512 && kUnitElems == 8); // Raw buffer word 3: OOB_SELECT 3, format 4. Bit 24 is reserved on gfx11 // and enables write compression on gfx12. static constexpr unsigned kRsrcWord3 = kGfx12 ? 0x30004000u : 0x31004000u; @@ -121,14 +135,30 @@ struct Prefill { largest_divisor_up_to(kNumKvHeads, C.head_group); static_assert(kNumQHeads % kNumKvHeads == 0, "GQA must be integral"); + static_assert(kCausal || kWindow == 0, "a non-causal build has no window"); + static_assert(kRowTiles == 1 || (kRowTiles == 2 && !kCausal && + kDimParts == 1 && kKeyWaves == 1), + "two row tiles per wave: non-causal, one part, one key wave"); + // Not yet ported to gfx12's operand layout. + static_assert(kRowTiles == 1 || !kGfx12, "two row tiles per wave: gfx11"); + // compute2 reads V^T from the staging buffer. + static_assert(kRowTiles == 1 || kVT, "two row tiles per wave read V^T"); + static_assert(kCausal || !kVTr, "v_t 2 reads V from the paged cache"); static_assert(kWaves % (kDimParts * kKeyWaves) == 0, "whole row tiles"); - static_assert(kHeadDim % (16 * kDimParts) == 0, "whole chunks per part"); + static_assert(kHeadDimP % (16 * kDimParts) == 0, "whole chunks per part"); + static_assert(kHeadDimP == kHeadDim || (kDimParts == 1 && kVT), + "a padded head dim is one part with V transposed"); + static_assert(!kVTr || kHeadDimP / 8 % 4 == 0, + "v_t 2 loads whole 16-dim chunks of 8-element units"); static_assert(kSubTiles % kKeyWaves == 0, "whole sub-tiles per key wave"); - static_assert(kLaneDims % 2 == 0, "a lane's V slice is whole dwords"); + static_assert(kVT || kLaneDims % 2 == 0, + "a lane's V slice is whole dwords, or V is transposed"); static_assert(kKeyTile % 16 == 0 && kPageSize % 16 == 0, "a 16-key sub-tile sits inside one page"); - static_assert(kRowChunks % 4 == 0, + static_assert(!kCausal || kRowChunks % 4 == 0, "a wave's keys of one load sit in one sub-tile, so one page"); + static_assert(kCausal || kHeadDim % 8 == 0, + "non-causal: whole 8-element units"); using elem_t = typename Elem::type; using e16 = typename Elem::v16; @@ -136,6 +166,29 @@ struct Prefill { using op_t = typename Elem::op; using opw_t = typename Elem::opw; static constexpr int kOpWords = Elem::kOpWords; + typedef elem_t e4 __attribute__((ext_vector_type(4))); + // One load unit. + using eu = std::conditional_t; + using uu = std::conditional_t; + + // A buffer load's 8 elements as a unit (buffer loads are 8-element units). + static __device__ __forceinline__ eu as_unit(e8 x) { + if constexpr (kUnitElems == 8) + return x; + else + return __builtin_shufflevector(x, x, 0, 1, 2, 3); + } + // The operand of a chunk with 8 real dims at `row` (D=72's last), zeros + // after. + static __device__ __forceinline__ op_t padded_op(const elem_t* row, + int laneHalf) { +#if defined(__GFX12__) + return laneHalf ? op_t{} : *(const op_t*)row; +#else + return __builtin_shufflevector(*(const e8*)row, e8{}, 0, 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, 13, 14, 15); +#endif + } using e2 = typename Elem::v2; static __device__ __forceinline__ elem_t to_elem(float x) { @@ -237,7 +290,9 @@ struct Prefill { const int* __restrict__ mmRanges, int numMmRanges, const int4* __restrict__ plan, const int* __restrict__ queryStart, int numEntries, - int blockTableStride, int headGroup) { + int blockTableStride, int headGroup, + const elem_t* __restrict__ vIn, int qStride, int kStride, + int vStride) { __shared__ __attribute__(( aligned(16))) elem_t kTile[kBuffers][kKeyTile][kRowPitch]; __shared__ __attribute__(( @@ -273,11 +328,18 @@ struct Prefill { groupSlot = (size_t)entry.w * kNumKvHeads + kvHead * numSegments; const int queryBegin = queryStart[seq]; numTokens = queryStart[seq + 1] - queryBegin; - q += (size_t)queryBegin * kNumQHeads * kHeadDim; out += (size_t)queryBegin * kNumQHeads * kHeadDim; - blockTable += (size_t)seq * blockTableStride; + if constexpr (kCausal) { + q += (size_t)queryBegin * kNumQHeads * kHeadDim; + blockTable += (size_t)seq * blockTableStride; + mmRanges += (size_t)seq * 2 * numMmRanges; + } else { + // The sequence's K and V start with its queries. + q += (size_t)queryBegin * qStride; + kv += (size_t)queryBegin * kStride; + vIn += (size_t)queryBegin * vStride; + } seqLens += seq; - mmRanges += (size_t)seq * 2 * numMmRanges; } else { // A tile's segments are adjacent; the heaviest tiles (last tokens) go // first, interleaving kHeadGroup kv heads. @@ -301,6 +363,7 @@ struct Prefill { // The last key any of tokens first..last sees: the last token, or the end // of a multimodal range they reach into. auto last_key = [&](int first, int last) __attribute__((always_inline)) { + if constexpr (!kCausal) return seqLen - 1; int key = context + last; if (numMmRanges == 0) return key; for (int r = 0; r < numMmRanges; ++r) { @@ -323,7 +386,7 @@ struct Prefill { numSegments * kKeyTile); const int lastPage = (seqLen - 1) / kPageSize; - const int waveRow = tileRow + rowWave * 16; + const int waveRow = tileRow + rowWave * 16 * kRowTiles; const int row = min(waveRow + laneRow, numRows - 1); const int rowToken = row / kG; const int rowHead = kvHead * kG + row % kG; @@ -339,15 +402,18 @@ struct Prefill { 1, kBlockRows * kHeadDim * (int)sizeof(elem_t) * 3 / (2 * 8192)); const int place = __builtin_amdgcn_readfirstlane((int)blockIdx.x); // Only when tile lengths differ and the KV is not split. - const bool unequal = numSegments == 1 && __builtin_amdgcn_readfirstlane( - *seqLens) <= 2 * numTokens; + const bool unequal = + kCausal && numSegments == 1 && + __builtin_amdgcn_readfirstlane(*seqLens) <= 2 * numTokens; if (unequal && place < firstRound) for (int i = 0; i < place * kStaggerUnits; ++i) __builtin_amdgcn_s_sleep(1); } // Lane l holds page l % kStepPages of a step, fetched a step ahead. + // Non-causal builds read K and V in place, without pages. auto step_pages = [&](int keyStart) __attribute__((always_inline)) { + if constexpr (!kCausal) return 0; return blockTable[min(keyStart / kPageSize + lane % kStepPages, lastPage)]; }; @@ -355,15 +421,50 @@ struct Prefill { op_t qOperand[kPartChunks]; #pragma unroll - for (int c = 0; c < kPartChunks; ++c) - qOperand[c] = Elem::load_op( - q + ((size_t)rowToken * kNumQHeads + rowHead) * kHeadDim + - (firstChunk + c) * 16, - laneHalf); + for (int c = 0; c < kPartChunks; ++c) { + const elem_t* qRow = + (kCausal ? q + ((size_t)rowToken * kNumQHeads + rowHead) * kHeadDim + : q + (size_t)rowToken * qStride + rowHead * kHeadDim) + + (firstChunk + c) * 16; + if (kHeadDimP != kHeadDim && (firstChunk + c) * 16 + 16 > kHeadDim) { + // The padded chunk: 8 real dims (D=72), zeros after. + qOperand[c] = padded_op(qRow, laneHalf); + } else { + qOperand[c] = Elem::load_op(qRow, laneHalf); + } + } + + // The second row tile's rows and Q (kRowTiles == 2). + const int row1 = min(waveRow + 16 + laneRow, numRows - 1); + op_t qOperand1[kRowTiles > 1 ? kPartChunks : 1]; + if constexpr (kRowTiles > 1) { +#pragma unroll + for (int c = 0; c < kPartChunks; ++c) { + const elem_t* qRow = q + (size_t)(row1 / kG) * qStride + + (kvHead * kG + row1 % kG) * kHeadDim + c * 16; + qOperand1[c] = kHeadDimP != kHeadDim && c * 16 + 16 > kHeadDim + ? padded_op(qRow, laneHalf) + : Elem::load_op(qRow, laneHalf); + } + } + + if constexpr (!kCausal && kHeadDimP != kHeadDim) { + // Staging never writes the padding: zero it once. + static_assert(kHeadDimP - kHeadDim == kUnitElems); + for (int i = thread; i < kBuffers * kKeyTile; i += kThreads) + *(eu*)&kTile[i / kKeyTile][i % kKeyTile][kHeadDim] = eu{}; + for (int i = thread; i < kBuffers * kUnitElems * kKeyTile; i += kThreads) + vTile[i / (kUnitElems * kKeyTile)] + [vt_row(kHeadDim + i / kKeyTile % kUnitElems)][i % kKeyTile] = + elem_t{}; + } - const elem_t* kvHeadBase = kv + (size_t)kvHead * (kPageSize * kKvRowElems); - e8 kStaged[kKLoads]; - e8 vStaged[kVLoads][kVPerUnit]; + const elem_t* kvHeadBase = + kCausal ? kv + (size_t)kvHead * (kPageSize * kKvRowElems) + : kv + kvHead * kHeadDim; + const elem_t* vHeadBase = kCausal ? nullptr : vIn + kvHead * kHeadDim; + eu kStaged[kKLoads]; + eu vStaged[kVLoads][kVPerUnit]; auto v_unit_pair = [&](int unit) __attribute__((always_inline)) { if constexpr (kVT) return unit / (8 * kRowChunks) * 8 + unit % 8; @@ -372,9 +473,9 @@ struct Prefill { }; auto v_unit_col = [&](int unit) __attribute__((always_inline)) { if constexpr (kVT) - return unit / 8 % kRowChunks * 8; + return unit / 8 % kRowChunks * kUnitElems; else - return unit % kRowChunks * 8; + return unit % kRowChunks * kUnitElems; }; auto v_unit_key = [&](int unit) __attribute__((always_inline)) { const int pair = v_unit_pair(unit); @@ -399,7 +500,7 @@ struct Prefill { kvHeadBase + (size_t)__builtin_amdgcn_readlane(pages, slot) * kPageElems + (first - slot * kPageSize) * kKvRowElems; - return *(const e8*)(base + (unsigned)(key * kKvRowElems + col)); + return *(const eu*)(base + (unsigned)(key * kKvRowElems + col)); } const int position = min(keyStart + key, seqLen - 1); const int slot = __builtin_amdgcn_readfirstlane(position / kPageSize - @@ -408,12 +509,33 @@ struct Prefill { kvHeadBase + (size_t)__builtin_amdgcn_readlane(pages, slot) * kPageElems; return *( - const e8*)(page + + const eu*)(page + (unsigned)((position % kPageSize) * kKvRowElems + col)); }; - // Per-lane byte offsets of the buffer loads, kept in registers. + // Per-lane byte offsets of the buffer loads, kept in registers: from the + // wave's first key, or (non-causal) the step's. int kOff[kKLoads], vOff[kVLoads][kVPerUnit]; - if constexpr (kBufferLoads) { + if constexpr (!kCausal) { +#pragma unroll + for (int i = 0; i < kKLoads; ++i) { + const int unit = i * kThreads + thread; + kOff[i] = + (unit / kRowChunks * kStride + unit % kRowChunks * kUnitElems) * + (int)sizeof(elem_t); + asm volatile("" : "+v"(kOff[i])); + } +#pragma unroll + for (int i = 0; i < kVLoads; ++i) { + const int unit = i * kThreads + thread; +#pragma unroll + for (int x = 0; x < kVPerUnit; ++x) { + vOff[i][x] = + ((v_unit_key(unit) + 2 * x) * vStride + v_unit_col(unit)) * + (int)sizeof(elem_t); + asm volatile("" : "+v"(vOff[i][x])); + } + } + } else if constexpr (kBufferLoads) { #pragma unroll for (int i = 0; i < kKLoads; ++i) { const int unit = i * kThreads + thread; @@ -450,6 +572,30 @@ struct Prefill { return __builtin_bit_cast( e8, __builtin_amdgcn_raw_buffer_load_b128(rsrc, offset, 0, 0)); }; + // Non-causal: the step's units from its first key, each range ending + // with its last key's row so later keys load as zeros. + auto issue_noncausal = [&](int keyStart) __attribute__((always_inline)) { + const int last = min(seqLen - keyStart, kKeyTile) - 1; + const __amdgpu_buffer_rsrc_t kRsrc = __builtin_amdgcn_make_buffer_rsrc( + (void*)(kvHeadBase + (size_t)keyStart * kStride), 0, + (last * kStride + kHeadDim) * (int)sizeof(elem_t), kRsrcWord3); + const __amdgpu_buffer_rsrc_t vRsrc = __builtin_amdgcn_make_buffer_rsrc( + (void*)(vHeadBase + (size_t)keyStart * vStride), 0, + (last * vStride + kHeadDim) * (int)sizeof(elem_t), kRsrcWord3); +#pragma unroll + for (int i = 0; i < kKLoads; ++i) + if (kKUnits % kThreads == 0 || i * kThreads + thread < kKUnits) + kStaged[i] = as_unit(__builtin_bit_cast( + e8, __builtin_amdgcn_raw_buffer_load_b128(kRsrc, kOff[i], 0, 0))); +#pragma unroll + for (int i = 0; i < kVLoads; ++i) + if (kVUnits % kThreads == 0 || i * kThreads + thread < kVUnits) +#pragma unroll + for (int x = 0; x < kVPerUnit; ++x) + vStaged[i][x] = as_unit( + __builtin_bit_cast(e8, __builtin_amdgcn_raw_buffer_load_b128( + vRsrc, vOff[i][x], 0, 0))); + }; // Load the step at keyStart into registers. auto issue_step = [&](int keyStart) __attribute__((always_inline)) { @@ -459,17 +605,17 @@ struct Prefill { #pragma unroll for (int i = 0; i < kKLoads; ++i) if (kKUnits % kThreads == 0 || i * kThreads + thread < kKUnits) - kStaged[i] = load_buffer( + kStaged[i] = as_unit(load_buffer( pages, keyStart, - (i * kThreads + wave * kWaveSize) / kRowChunks, kOff[i]); + (i * kThreads + wave * kWaveSize) / kRowChunks, kOff[i])); #pragma unroll for (int i = 0; i < kVIssue; ++i) if (kVUnits % kThreads == 0 || i * kThreads + thread < kVUnits) { const int waveKey = v_unit_key(i * kThreads + wave * kWaveSize); #pragma unroll for (int x = 0; x < kVPerUnit; ++x) - vStaged[i][x] = - load_buffer(pages, keyStart, waveKey + 2 * x, vOff[i][x]); + vStaged[i][x] = as_unit(load_buffer( + pages, keyStart, waveKey + 2 * x, vOff[i][x])); } return; } @@ -478,7 +624,8 @@ struct Prefill { const int unit = i * kThreads + thread; if (kKUnits % kThreads == 0 || unit < kKUnits) kStaged[i] = load_row.template operator()( - pages, keyStart, unit / kRowChunks, unit % kRowChunks * 8); + pages, keyStart, unit / kRowChunks, + min(unit % kRowChunks * kUnitElems, kHeadDim - kUnitElems)); } #pragma unroll for (int i = 0; i < kVIssue; ++i) { @@ -488,7 +635,8 @@ struct Prefill { #pragma unroll for (int x = 0; x < kVPerUnit; ++x) vStaged[i][x] = load_row.template operator()( - pages, keyStart, key + 2 * x, kHeadDim + col); + pages, keyStart, key + 2 * x, + kHeadDim + min(col, kHeadDim - kUnitElems)); } } }; @@ -503,9 +651,11 @@ struct Prefill { for (int i = 0; i < (kDoK ? kKLoads : 0); ++i) { const int unit = i * kThreads + thread; if (kKUnits % kThreads == 0 || unit < kKUnits) { - const int key = unit / kRowChunks, col = unit % kRowChunks * 8; - *(e8*)&kTile[buffer][key][col] = - !kEdge || valid(keyStart + key) ? kStaged[i] : e8{}; + const int key = unit / kRowChunks, + col = unit % kRowChunks * kUnitElems; + *(eu*)&kTile[buffer][key][col] = + (!kEdge || valid(keyStart + key)) && col < kHeadDim ? kStaged[i] + : eu{}; } } #pragma unroll @@ -514,46 +664,52 @@ struct Prefill { if (kVUnits % kThreads == 0 || unit < kVUnits) { const int key = v_unit_key(unit), col = v_unit_col(unit); if constexpr (kVT) { - const e8 a = !kEdge || valid(keyStart + key) ? vStaged[i][0] : e8{}; - const e8 b = - !kEdge || valid(keyStart + key + 2) ? vStaged[i][1] : e8{}; - const u4v aw = __builtin_bit_cast(u4v, a), - bw = __builtin_bit_cast(u4v, b); + const eu a = (!kEdge || valid(keyStart + key)) && col < kHeadDim + ? vStaged[i][0] + : eu{}; + const eu b = (!kEdge || valid(keyStart + key + 2)) && col < kHeadDim + ? vStaged[i][1] + : eu{}; + const uu aw = __builtin_bit_cast(uu, a), + bw = __builtin_bit_cast(uu, b); const int slot = key / 16 * 16 + vt_slot(key % 16); #pragma unroll - for (int d = 0; d < 8; ++d) + for (int d = 0; d < kUnitElems; ++d) *(unsigned*)&vTile[buffer][vt_row(col + d)][slot] = __builtin_amdgcn_perm(bw[d >> 1], aw[d >> 1], (d & 1) ? 0x07060302u : 0x05040100u); } else { - *(e8*)&vTile[buffer][key][col] = - !kEdge || valid(keyStart + key) ? vStaged[i][0] : e8{}; + *(eu*)&vTile[buffer][key][col] = + !kEdge || valid(keyStart + key) ? vStaged[i][0] : eu{}; } } } }; auto issue = [&](int keyStart) __attribute__((always_inline)) { - if (edge_step(keyStart)) + if constexpr (!kCausal) + issue_noncausal(keyStart); + else if (edge_step(keyStart)) issue_step.template operator()(keyStart); else issue_step.template operator()(keyStart); }; + // Non-causal steps past S loaded zeros: no masking. auto stage = [&](int buffer, int keyStart) __attribute__((always_inline)) { - if (edge_step(keyStart)) + if (kCausal && edge_step(keyStart)) stage_step.template operator()(buffer, keyStart); else stage_step.template operator()(buffer, keyStart); }; auto stage_k = [&](int buffer, int keyStart) __attribute__((always_inline)) { - if (edge_step(keyStart)) + if (kCausal && edge_step(keyStart)) stage_step.template operator()(buffer, keyStart); else stage_step.template operator()(buffer, keyStart); }; auto stage_v = [&](int buffer, int keyStart) __attribute__((always_inline)) { - if (edge_step(keyStart)) + if (kCausal && edge_step(keyStart)) stage_step.template operator()(buffer, keyStart); else stage_step.template operator()(buffer, keyStart); @@ -567,11 +723,103 @@ struct Prefill { f8 accum[kLaneDims]; #pragma unroll for (int j = 0; j < kLaneDims; ++j) accum[j] = f8{}; + float runMax1 = -INFINITY, runSum1 = 0.f; + f8 accum1[kRowTiles > 1 ? kLaneDims : 1]; +#pragma unroll + for (int j = 0; j < (kRowTiles > 1 ? kLaneDims : 1); ++j) accum1[j] = f8{}; // Q@K, softmax update and P@V for the step at keyStart in `buffer`. Lane // holds row laneRow, keys keyStart + 16 st + 2e + laneHalf. + // Online softmax of one row tile's scores: rescales its O, returns P. + auto softmax = [&](f8* scores, float& rowMax, float& rowSum, f8* acc, + op_t* p) __attribute__((always_inline)) { + float tileMax = -INFINITY; +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) +#pragma unroll + for (int e = 0; e < 8; ++e) tileMax = fmaxf(tileMax, scores[ls][e]); + tileMax = fmaxf(tileMax, xhalf(tileMax)) * scaleLog2; + const float newMax = tileMax > rowMax + kMaxSlack ? tileMax : rowMax; + const float safeMax = newMax == -INFINITY ? 0.f : newMax; + const float rescale = __builtin_amdgcn_exp2f(rowMax - safeMax); + rowMax = newMax; + if (__builtin_amdgcn_ballot_w32(rescale != 1.f)) +#pragma unroll + for (int j = 0; j < kLaneDims; ++j) acc[j] *= rescale; + float tileSum = 0.f; +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) { + float probs[8]; +#pragma unroll + for (int e = 0; e < 8; ++e) { + probs[e] = __builtin_amdgcn_exp2f( + __builtin_fmaf(scores[ls][e], scaleLog2, -safeMax)); + tileSum += probs[e]; + } + p[ls] = p_frag(probs); + } + rowSum = rowSum * rescale + tileSum + xhalf(tileSum); + }; + // Two row tiles (non-causal): each K and V operand read feeds both. + auto compute2 = [&](int buffer, + int keyStart) __attribute__((always_inline)) { + if constexpr (kRowTiles > 1) { + bool live[kWaveSubTiles]; + bool allLive = true; +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) { + live[ls] = keyStart + ls * 16 < seqLen; + allLive = allLive && live[ls]; + } + allLive = __builtin_amdgcn_readfirstlane((int)allLive); + f8 s0[kWaveSubTiles], s1[kWaveSubTiles]; +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) s0[ls] = s1[ls] = f8{}; +#pragma unroll + for (int c = 0; c < kPartChunks; ++c) +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) { + const e16 kOperand = + *(const e16*)&kTile[buffer][ls * 16 + laneRow][c * 16]; + if (allLive || live[ls]) { + s0[ls] = wmma(kOperand, qOperand[c], s0[ls]); + s1[ls] = wmma(kOperand, qOperand1[c], s1[ls]); + } + } + if (__builtin_expect(keyStart + kKeyTile > seqLen, 0)) { +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) +#pragma unroll + for (int e = 0; e < 8; ++e) + if (keyStart + ls * 16 + 2 * e + laneHalf >= seqLen) + s0[ls][e] = s1[ls][e] = -INFINITY; + } + op_t p0[kWaveSubTiles], p1[kWaveSubTiles]; + softmax(s0, runMax, runSum, accum, p0); + softmax(s1, runMax1, runSum1, accum1, p1); +#pragma unroll + for (int ls = 0; ls < kWaveSubTiles; ++ls) +#pragma unroll + for (int j = 0; j < kLaneDims; ++j) { + const elem_t* vRow = &vTile[buffer][j * 16 + laneRow][ls * 16]; + const e8 own = *(const e8*)(vRow + laneHalf * 8); + const e8 other = *(const e8*)(vRow + (1 - laneHalf) * 8); + const e16 vOperand = + __builtin_shufflevector(own, other, 0, 1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15); + if (allLive || live[ls]) { + accum[j] = wmma(vOperand, p0[ls], accum[j]); + accum1[j] = wmma(vOperand, p1[ls], accum1[j]); + } + } + } + }; auto compute = [&](int buffer, int keyStart, int stageNext = -1) __attribute__((always_inline)) { + if constexpr (kRowTiles > 1) { + compute2(buffer, keyStart); + return; + } // v_t 2: the step's pages, read now so Q@K hides the latency. const int vPages = kVTr ? step_pages(keyStart) : 0; // The global address of key `first` + k of a 16-key sub-tile. @@ -710,8 +958,10 @@ struct Prefill { } // Mask only diagonal and window-edge steps. No row's last key is before - // the first token's. - const bool causalEdge = keyStart + kKeyTile - 1 > context + firstToken; + // the first token's. Non-causal: only the step past the last key. + const bool causalEdge = + kCausal ? keyStart + kKeyTile - 1 > context + firstToken + : keyStart + kKeyTile > seqLen; const bool windowEdge = kWindow && keyStart < context + lastToken - (kWindow - 1); if (__builtin_expect(causalEdge || windowEdge, 0)) { @@ -907,6 +1157,41 @@ struct Prefill { } const bool owner = kKeyWaves == 1 || keyWave == 0; + if constexpr (kRowTiles > 1) { + // The second row tile first: the first's store may return. + if (numSegments == 1) { + if (waveRow + 16 + laneRow < numRows) { + const float inv = __builtin_amdgcn_rcpf(runSum1); + elem_t* dst = out + ((size_t)(row1 / kG) * kNumQHeads + kvHead * kG + + row1 % kG) * + kHeadDim; +#pragma unroll + for (int e = 0; e < 8; ++e) +#pragma unroll + for (int j = 0; j < kLaneDims; ++j) + if (kHeadDimP == kHeadDim || + kLaneDims * acc_row(e, laneHalf) + j < kHeadDim) + dst[kLaneDims * acc_row(e, laneHalf) + j] = + to_elem(accum1[j][e] * inv); + } + } else { + const size_t partialRow = (groupSlot + segment) * kBlockRows + + rowWave * 16 * kRowTiles + 16 + laneRow; + float* dst = partialO + partialRow * kHeadDim; +#pragma unroll + for (int e = 0; e < 8; ++e) +#pragma unroll + for (int j = 0; j < kLaneDims; ++j) + if (kHeadDimP == kHeadDim || + kLaneDims * acc_row(e, laneHalf) + j < kHeadDim) + dst[kLaneDims * acc_row(e, laneHalf) + j] = accum1[j][e]; + if (laneHalf == 0) { + partialML[partialRow * 2] = runMax1; + partialML[partialRow * 2 + 1] = runSum1; + } + } + } + if (numSegments == 1) { if (!owner || waveRow + laneRow >= numRows) return; const float inv = __builtin_amdgcn_rcpf(runSum); @@ -928,10 +1213,18 @@ struct Prefill { elem_t x[kLaneDims]; #pragma unroll for (int j = 0; j < kLaneDims; ++j) x[j] = to_elem(accum[j][e] * inv); - unsigned* d = (unsigned*)(dst + kLaneDims * acc_row(e, laneHalf)); + if constexpr (kLaneDims % 2) { +#pragma unroll + for (int j = 0; j < kLaneDims; ++j) + if (kHeadDimP == kHeadDim || + kLaneDims * acc_row(e, laneHalf) + j < kHeadDim) + dst[kLaneDims * acc_row(e, laneHalf) + j] = x[j]; + } else { + unsigned* d = (unsigned*)(dst + kLaneDims * acc_row(e, laneHalf)); #pragma unroll - for (int w = 0; w < kLaneWords; ++w) - d[w] = __builtin_bit_cast(unsigned, e2{x[2 * w], x[2 * w + 1]}); + for (int w = 0; w < kLaneWords; ++w) + d[w] = __builtin_bit_cast(unsigned, e2{x[2 * w], x[2 * w + 1]}); + } } return; } @@ -940,14 +1233,17 @@ struct Prefill { // resets the counter. const size_t partialBase = (groupSlot + segment) * kBlockRows; if (owner) { - const size_t partialRow = partialBase + rowWave * 16 + laneRow; + const size_t partialRow = + partialBase + rowWave * 16 * kRowTiles + laneRow; float* dst = partialO + partialRow * kHeadDim + dimPart * kPartDims; #pragma unroll for (int e = 0; e < 8; ++e) #pragma unroll for (int j = 0; j < kLaneDims; ++j) - dst[kVTr ? 16 * j + acc_row(e, laneHalf) - : kLaneDims * acc_row(e, laneHalf) + j] = accum[j][e]; + if (kHeadDimP == kHeadDim || + kLaneDims * acc_row(e, laneHalf) + j < kHeadDim) + dst[kVTr ? 16 * j + acc_row(e, laneHalf) + : kLaneDims * acc_row(e, laneHalf) + j] = accum[j][e]; if (laneHalf == 0 && dimPart == 0) { partialML[partialRow * 2] = runMax; partialML[partialRow * 2 + 1] = runSum; @@ -1047,7 +1343,8 @@ void launch_prefill(const PrefillArgs& a) { a.seq_lens, 0, a.scale, a.partial_o, a.partial_ml, a.counters, 1, wgpWorkgroups * numWgps, a.mm_ranges, a.num_mm_ranges, reinterpret_cast(a.plan), a.query_start, a.num_entries, - a.bt_stride, headGroup); + a.bt_stride, headGroup, static_cast(a.v), a.q_stride, + a.k_stride, a.v_stride); return; } // A workgroup per SIMD. @@ -1073,7 +1370,8 @@ void launch_prefill(const PrefillArgs& a) { static_cast(a.kv), a.bt, static_cast(a.out), a.seq_lens, a.num_tokens, a.scale, a.partial_o, a.partial_ml, a.counters, segments, wgpWorkgroups * numWgps, a.mm_ranges, a.num_mm_ranges, nullptr, - nullptr, 0, 0, P::kHeadGroup); + nullptr, 0, 0, P::kHeadGroup, static_cast(a.v), a.q_stride, + a.k_stride, a.v_stride); } } // namespace diff --git a/csrc/rocm/rdna_prefill_batch_configs.h b/csrc/rocm/rdna_prefill_batch_configs.h index 4137cd68971e..341f9c0a2246 100644 --- a/csrc/rocm/rdna_prefill_batch_configs.h +++ b/csrc/rocm/rdna_prefill_batch_configs.h @@ -18,6 +18,7 @@ struct PrefillBatchConfig { int min_keys; int min_query_len; // of the kPrefillConfigs entry, in both dtypes int head_group; // 0: the build's + int causal = 1; // of the kPrefillConfigs entry; left out when 1 }; // One table per Arch (rdna_attn_types.h), each complete; a GPU family's @@ -26,6 +27,87 @@ struct PrefillBatchConfig { // line in PrefillBatchConfig order. // clang-format off inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { + {HeadDim::k32, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 16, 0, 32, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 64, 0, 1536, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 0, 16, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 512, 32, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 512, 1536, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 2048, 1536, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 2048, 1536, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 8192, 1536, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 16, 0, 512, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 64, 0, 128, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 0, 32, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 512, 32, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 512, 1536, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 2048, 1536, 1, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 8192, 1536, 1, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 64, 0, 1536, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 0, 1024, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 512, 1024, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 512, 1024, 2, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 2048, 1024, 8, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 2048, 1024, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 8192, 1024, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1, 0, 16, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 16, 0, 32, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 64, 0, 32, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 0, 1024, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 512, 1536, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 512, 3072, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 2048, 1024, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 2048, 1024, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 8192, 1024, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1, 0, 16, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 16, 0, 16, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 64, 0, 768, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 0, 4096, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 512, 4096, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 512, 4096, 1, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 2048, 4096, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 2048, 4096, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 8192, 4096, 8, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 16, 0, 64, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 64, 0, 3072, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 0, 64, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 512, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 512, 3072, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 2048, 3072, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 2048, 3072, 1, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 8192, 3072, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1, 0, 1024, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 16, 0, 16, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 64, 0, 64, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 0, 1536, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 512, 1536, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 512, 64, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 8192, 1536, 3, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1, 0, 32, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 16, 0, 768, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 64, 0, 32, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 0, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 512, 1536, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 512, 16, 1, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 2048, 8192, 1, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 8192, 1536, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1, 0, 16, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 16, 0, 256, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 64, 0, 256, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 0, 1536, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 512, 1024, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 512, 1536, 2, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 2048, 1536, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 8192, 1536, 1, 0}, {HeadDim::k64, 14, 2, 0, 16, 1, 0, 16, 2}, {HeadDim::k64, 14, 2, 0, 16, 1, 512, 16, 2}, {HeadDim::k64, 14, 2, 0, 16, 1, 2048, 16, 2}, @@ -47,6 +129,15 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k64, 14, 2, 0, 16, 1024, 8192, 128, 1}, {HeadDim::k64, 14, 2, 0, 16, 4096, 2048, 128, 1}, {HeadDim::k64, 14, 2, 0, 16, 4096, 8192, 128, 1}, + {HeadDim::k64, 14, 14, 0, 16, 1, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 16, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 64, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 0, 768, 7, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 512, 8192, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 512, 1536, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 2048, 8192, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 8192, 8192, 1, 0}, {HeadDim::k64, 16, 2, 0, 16, 1, 0, 16, 2}, {HeadDim::k64, 16, 2, 0, 16, 1, 512, 16, 2}, {HeadDim::k64, 16, 2, 0, 16, 1, 2048, 16, 1}, @@ -68,6 +159,33 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k64, 16, 2, 0, 16, 1024, 8192, 128, 1}, {HeadDim::k64, 16, 2, 0, 16, 4096, 2048, 128, 1}, {HeadDim::k64, 16, 2, 0, 16, 4096, 8192, 128, 1}, + {HeadDim::k64, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 64, 0, 768, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 0, 768, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 512, 4096, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 512, 4096, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 2048, 4096, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 2048, 4096, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 8192, 768, 8, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1, 0, 16, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 16, 0, 16, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 64, 0, 1536, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 0, 512, 10, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 512, 512, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 512, 512, 5, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 2048, 1536, 5, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 8192, 1536, 1, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1, 0, 32, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 16, 0, 768, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 64, 0, 768, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 0, 512, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 512, 1536, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 512, 1536, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 2048, 1536, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 8192, 1536, 1, 0}, {HeadDim::k64, 32, 8, 0, 16, 1, 0, 16, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 512, 16, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 2048, 16, 8}, @@ -110,6 +228,33 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k64, 32, 32, 0, 16, 1024, 8192, 512, 1}, {HeadDim::k64, 32, 32, 0, 16, 4096, 2048, 512, 1}, {HeadDim::k64, 32, 32, 0, 16, 4096, 8192, 512, 1}, + {HeadDim::k64, 64, 64, 0, 16, 1, 0, 16, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 16, 0, 32, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 64, 0, 32, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 0, 16, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 512, 256, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 512, 256, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 2048, 256, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 2048, 256, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 8192, 8192, 32, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 16, 0, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 64, 0, 1536, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 0, 768, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 512, 768, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 512, 768, 4, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 2048, 768, 1, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 2048, 768, 1, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 8192, 768, 1, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 16, 0, 16, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 64, 0, 1536, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 0, 768, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 512, 768, 2, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 512, 768, 1, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 2048, 768, 1, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 2048, 8192, 2, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 8192, 8192, 2, 0}, {HeadDim::k96, 32, 32, 0, 16, 1, 0, 64, 32}, {HeadDim::k96, 32, 32, 0, 16, 1, 512, 64, 32}, {HeadDim::k96, 32, 32, 0, 16, 1, 2048, 16, 32}, @@ -131,6 +276,24 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k96, 32, 32, 0, 16, 1024, 8192, 2048, 1}, {HeadDim::k96, 32, 32, 0, 16, 4096, 2048, 2048, 1}, {HeadDim::k96, 32, 32, 0, 16, 4096, 8192, 2048, 1}, + {HeadDim::k128, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 16, 0, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 64, 0, 32, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 0, 768, 2, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 512, 3072, 1, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 512, 1024, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 2048, 512, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 2048, 1024, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 8192, 3072, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 16, 0, 32, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 64, 0, 256, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 0, 16, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 512, 3072, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 512, 3072, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 2048, 3072, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 2048, 3072, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 8192, 1024, 4, 0}, {HeadDim::k128, 10, 10, 0, 16, 1, 0, 16, 10}, {HeadDim::k128, 10, 10, 0, 16, 1, 512, 16, 2}, {HeadDim::k128, 10, 10, 0, 16, 1, 2048, 16, 10}, @@ -152,6 +315,15 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k128, 10, 10, 0, 16, 1024, 8192, 512, 1}, {HeadDim::k128, 10, 10, 0, 16, 4096, 2048, 512, 1}, {HeadDim::k128, 10, 10, 0, 16, 4096, 8192, 512, 2}, + {HeadDim::k128, 12, 12, 0, 16, 1, 0, 768, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 16, 0, 768, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 64, 0, 4096, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 0, 256, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 512, 4096, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 512, 4096, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 2048, 4096, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 2048, 4096, 3, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 8192, 4096, 1, 0}, {HeadDim::k128, 16, 2, 0, 16, 1, 0, 16, 2}, {HeadDim::k128, 16, 2, 0, 16, 1, 512, 16, 2}, {HeadDim::k128, 16, 2, 0, 16, 1, 2048, 16, 2}, @@ -194,6 +366,15 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k128, 16, 8, 0, 16, 1024, 8192, 256, 2}, {HeadDim::k128, 16, 8, 0, 16, 4096, 2048, 256, 2}, {HeadDim::k128, 16, 8, 0, 16, 4096, 8192, 256, 1}, + {HeadDim::k128, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 64, 0, 32, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 0, 3072, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 512, 3072, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 512, 8192, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 2048, 3072, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 2048, 1024, 8, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 8192, 1024, 4, 0}, {HeadDim::k128, 24, 8, 0, 16, 1, 0, 16, 8}, {HeadDim::k128, 24, 8, 0, 16, 1, 512, 16, 1}, {HeadDim::k128, 24, 8, 0, 16, 1, 2048, 16, 8}, @@ -215,6 +396,15 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { {HeadDim::k128, 24, 8, 0, 16, 1024, 8192, 512, 1}, {HeadDim::k128, 24, 8, 0, 16, 4096, 2048, 512, 1}, {HeadDim::k128, 24, 8, 0, 16, 4096, 8192, 512, 1}, + {HeadDim::k128, 24, 24, 0, 16, 1, 0, 16, 24, 0}, + {HeadDim::k128, 24, 24, 0, 16, 16, 0, 16, 24, 0}, + {HeadDim::k128, 24, 24, 0, 16, 64, 0, 16, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 0, 256, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 512, 256, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 512, 256, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 2048, 256, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 2048, 1536, 3, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 8192, 1536, 2, 0}, {HeadDim::k128, 28, 4, 0, 16, 1, 0, 16, 4}, {HeadDim::k128, 28, 4, 0, 16, 1, 512, 16, 4}, {HeadDim::k128, 28, 4, 0, 16, 1, 2048, 16, 1}, @@ -828,6 +1018,87 @@ inline constexpr PrefillBatchConfig kGfx1151PrefillBatchConfigs[] = { }; inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { + {HeadDim::k32, 4, 4, 0, 16, 1, 0, 3072, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 16, 0, 512, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 64, 0, 1536, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 0, 1536, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 512, 1536, 1, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 512, 1024, 1, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 2048, 1024, 1, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 2048, 1024, 1, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 8192, 1024, 4, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 16, 0, 32, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 64, 0, 1024, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 0, 768, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 512, 1024, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 512, 1024, 1, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 2048, 1024, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 2048, 1024, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 8192, 768, 8, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1, 0, 32, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 16, 0, 16, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 64, 0, 512, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 0, 512, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 512, 4096, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 512, 512, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 2048, 512, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 2048, 8192, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 8192, 8192, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 16, 0, 16, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 64, 0, 3072, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 0, 3072, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 512, 3072, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 512, 3072, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 2048, 3072, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 2048, 4096, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 8192, 4096, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1, 0, 32, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 16, 0, 16, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 64, 0, 128, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 0, 128, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 512, 1024, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 512, 1024, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 2048, 128, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 2048, 1024, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 8192, 1024, 16, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 16, 0, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 64, 0, 512, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 0, 768, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 512, 128, 1, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 512, 768, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 2048, 768, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 2048, 768, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 8192, 768, 4, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1, 0, 128, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 16, 0, 32, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 64, 0, 128, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 0, 32, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 512, 32, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 512, 32, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 2048, 512, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 2048, 512, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 8192, 512, 1, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 16, 0, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 64, 0, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 0, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 512, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 512, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 2048, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 2048, 8192, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 8192, 8192, 8, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1, 0, 32, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 16, 0, 32, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 64, 0, 256, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 0, 64, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 512, 64, 3, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 512, 1536, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 2048, 1536, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 2048, 1536, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 8192, 1536, 4, 0}, {HeadDim::k64, 14, 2, 0, 16, 1, 0, 512, 2}, {HeadDim::k64, 14, 2, 0, 16, 1, 512, 512, 1}, {HeadDim::k64, 14, 2, 0, 16, 1, 2048, 512, 1}, @@ -849,6 +1120,15 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k64, 14, 2, 0, 16, 1024, 8192, 64, 1}, {HeadDim::k64, 14, 2, 0, 16, 4096, 2048, 512, 2}, {HeadDim::k64, 14, 2, 0, 16, 4096, 8192, 64, 1}, + {HeadDim::k64, 14, 14, 0, 16, 1, 0, 32, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 16, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 64, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 0, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 512, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 512, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 2048, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 8192, 1536, 1, 0}, {HeadDim::k64, 16, 2, 0, 16, 1, 0, 2048, 1}, {HeadDim::k64, 16, 2, 0, 16, 1, 512, 16, 1}, {HeadDim::k64, 16, 2, 0, 16, 1, 2048, 16, 1}, @@ -870,6 +1150,33 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k64, 16, 2, 0, 16, 1024, 8192, 64, 1}, {HeadDim::k64, 16, 2, 0, 16, 4096, 2048, 2048, 2}, {HeadDim::k64, 16, 2, 0, 16, 4096, 8192, 64, 2}, + {HeadDim::k64, 16, 16, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 16, 0, 16, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 64, 0, 32, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 0, 256, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 512, 128, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 512, 128, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 2048, 128, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 2048, 8192, 8, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 8192, 8192, 8, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 16, 0, 16, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 64, 0, 1536, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 0, 256, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 512, 1536, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 512, 1536, 1, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 2048, 1536, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 2048, 1536, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 8192, 1536, 4, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1, 0, 128, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 16, 0, 128, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 64, 0, 512, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 0, 256, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 512, 8192, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 512, 512, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 2048, 512, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 2048, 8192, 12, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 8192, 8192, 6, 0}, {HeadDim::k64, 32, 8, 0, 16, 1, 0, 128, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 512, 128, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 2048, 16, 8}, @@ -912,6 +1219,33 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k64, 32, 32, 0, 16, 1024, 8192, 256, 1}, {HeadDim::k64, 32, 32, 0, 16, 4096, 2048, 256, 2}, {HeadDim::k64, 32, 32, 0, 16, 4096, 8192, 256, 1}, + {HeadDim::k64, 64, 64, 0, 16, 1, 0, 16, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 16, 0, 16, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 64, 0, 512, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 0, 512, 8, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 512, 512, 8, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 512, 512, 8, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 2048, 512, 8, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 2048, 3072, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 8192, 3072, 32, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1, 0, 256, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 16, 0, 64, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 64, 0, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 0, 64, 4, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 512, 512, 2, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 512, 1536, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 2048, 1536, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 2048, 512, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 8192, 1536, 1, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 16, 0, 16, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 64, 0, 64, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 0, 512, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 512, 1024, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 512, 8192, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 2048, 8192, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 2048, 1024, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 8192, 8192, 2, 0}, {HeadDim::k96, 32, 32, 0, 16, 1, 0, 16, 32}, {HeadDim::k96, 32, 32, 0, 16, 1, 512, 16, 32}, {HeadDim::k96, 32, 32, 0, 16, 1, 2048, 16, 32}, @@ -933,6 +1267,24 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k96, 32, 32, 0, 16, 1024, 8192, 1024, 1}, {HeadDim::k96, 32, 32, 0, 16, 4096, 2048, 1024, 2}, {HeadDim::k96, 32, 32, 0, 16, 4096, 8192, 1024, 1}, + {HeadDim::k128, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 16, 0, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 64, 0, 512, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 0, 32, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 512, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 512, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 2048, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 2048, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 8192, 768, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1, 0, 256, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 16, 0, 256, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 64, 0, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 0, 512, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 512, 512, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 512, 512, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 2048, 512, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 2048, 512, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 8192, 512, 4, 0}, {HeadDim::k128, 10, 10, 0, 16, 1, 0, 16, 10}, {HeadDim::k128, 10, 10, 0, 16, 1, 512, 16, 10}, {HeadDim::k128, 10, 10, 0, 16, 1, 2048, 16, 10}, @@ -954,6 +1306,15 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k128, 10, 10, 0, 16, 1024, 8192, 512, 1}, {HeadDim::k128, 10, 10, 0, 16, 4096, 2048, 512, 1}, {HeadDim::k128, 10, 10, 0, 16, 4096, 8192, 512, 1}, + {HeadDim::k128, 12, 12, 0, 16, 1, 0, 256, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 16, 0, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 64, 0, 768, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 0, 768, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 512, 768, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 512, 512, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 2048, 768, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 2048, 512, 4, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 8192, 768, 6, 0}, {HeadDim::k128, 16, 2, 0, 16, 1, 0, 128, 1}, {HeadDim::k128, 16, 2, 0, 16, 1, 512, 16, 2}, {HeadDim::k128, 16, 2, 0, 16, 1, 2048, 16, 2}, @@ -996,6 +1357,15 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k128, 16, 8, 0, 16, 1024, 8192, 128, 1}, {HeadDim::k128, 16, 8, 0, 16, 4096, 2048, 128, 1}, {HeadDim::k128, 16, 8, 0, 16, 4096, 8192, 128, 1}, + {HeadDim::k128, 16, 16, 0, 16, 1, 0, 32, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 16, 0, 32, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 64, 0, 32, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 0, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 512, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 512, 1536, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 2048, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 2048, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 8192, 16, 4, 0}, {HeadDim::k128, 24, 8, 0, 16, 1, 0, 16, 8}, {HeadDim::k128, 24, 8, 0, 16, 1, 512, 16, 8}, {HeadDim::k128, 24, 8, 0, 16, 1, 2048, 16, 8}, @@ -1017,6 +1387,15 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { {HeadDim::k128, 24, 8, 0, 16, 1024, 8192, 32, 1}, {HeadDim::k128, 24, 8, 0, 16, 4096, 2048, 512, 1}, {HeadDim::k128, 24, 8, 0, 16, 4096, 8192, 32, 1}, + {HeadDim::k128, 24, 24, 0, 16, 1, 0, 128, 12, 0}, + {HeadDim::k128, 24, 24, 0, 16, 16, 0, 16, 12, 0}, + {HeadDim::k128, 24, 24, 0, 16, 64, 0, 128, 8, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 0, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 512, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 512, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 2048, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 2048, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 8192, 128, 8, 0}, {HeadDim::k128, 28, 4, 0, 16, 1, 0, 16, 4}, {HeadDim::k128, 28, 4, 0, 16, 1, 512, 16, 2}, {HeadDim::k128, 28, 4, 0, 16, 1, 2048, 64, 4}, @@ -1630,6 +2009,87 @@ inline constexpr PrefillBatchConfig kGfx1150PrefillBatchConfigs[] = { }; inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { + {HeadDim::k32, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 16, 0, 64, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 64, 0, 8192, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 0, 32, 1, 0}, + {HeadDim::k32, 4, 4, 0, 16, 256, 512, 64, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 512, 64, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 1024, 2048, 8192, 4, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 2048, 64, 2, 0}, + {HeadDim::k32, 4, 4, 0, 16, 4096, 8192, 8192, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 16, 0, 64, 1, 0}, + {HeadDim::k32, 8, 8, 0, 16, 64, 0, 768, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 0, 32, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 256, 512, 32, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 512, 1536, 8, 0}, + {HeadDim::k32, 8, 8, 0, 16, 1024, 2048, 32, 1, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k32, 8, 8, 0, 16, 4096, 8192, 1536, 4, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 16, 0, 512, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 64, 0, 1536, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 0, 1536, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 256, 512, 1536, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 512, 1536, 16, 0}, + {HeadDim::k32, 16, 16, 0, 16, 1024, 2048, 1536, 8, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k32, 16, 16, 0, 16, 4096, 8192, 1536, 2, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1, 0, 16, 32, 0}, + {HeadDim::k32, 32, 32, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 64, 0, 512, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 0, 512, 2, 0}, + {HeadDim::k32, 32, 32, 0, 16, 256, 512, 512, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 512, 512, 16, 0}, + {HeadDim::k32, 32, 32, 0, 16, 1024, 2048, 512, 8, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 2048, 512, 2, 0}, + {HeadDim::k32, 32, 32, 0, 16, 4096, 8192, 512, 2, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1, 0, 16, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 16, 0, 32, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 64, 0, 512, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 0, 512, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 256, 512, 512, 32, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 512, 16, 1, 0}, + {HeadDim::k48, 32, 32, 0, 16, 1024, 2048, 512, 8, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 2048, 512, 4, 0}, + {HeadDim::k48, 32, 32, 0, 16, 4096, 8192, 512, 2, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 16, 0, 64, 2, 0}, + {HeadDim::k64, 4, 4, 0, 16, 64, 0, 768, 2, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 0, 768, 1, 0}, + {HeadDim::k64, 4, 4, 0, 16, 256, 512, 768, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 512, 1536, 4, 0}, + {HeadDim::k64, 4, 4, 0, 16, 1024, 2048, 768, 1, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k64, 4, 4, 0, 16, 4096, 8192, 1536, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1, 0, 16, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 16, 0, 64, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 64, 0, 256, 6, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 0, 256, 3, 0}, + {HeadDim::k64, 6, 6, 0, 16, 256, 512, 256, 3, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 512, 256, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 1024, 2048, 8192, 1, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k64, 6, 6, 0, 16, 4096, 8192, 1536, 2, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1, 0, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 16, 0, 64, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 64, 0, 64, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 0, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 256, 512, 8192, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 512, 1024, 8, 0}, + {HeadDim::k64, 8, 8, 0, 16, 1024, 2048, 8192, 1, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 2048, 8192, 1, 0}, + {HeadDim::k64, 8, 8, 0, 16, 4096, 8192, 8192, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1, 0, 16, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 16, 0, 16, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 64, 0, 16, 12, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 0, 128, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 256, 512, 4096, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 512, 4096, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 1024, 2048, 4096, 1, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 2048, 4096, 4, 0}, + {HeadDim::k64, 12, 12, 0, 16, 4096, 8192, 4096, 2, 0}, {HeadDim::k64, 14, 2, 0, 16, 1, 0, 32, 2}, {HeadDim::k64, 14, 2, 0, 16, 1, 512, 32, 2}, {HeadDim::k64, 14, 2, 0, 16, 1, 2048, 16, 2}, @@ -1651,6 +2111,15 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k64, 14, 2, 0, 16, 1024, 8192, 128, 2}, {HeadDim::k64, 14, 2, 0, 16, 4096, 2048, 128, 2}, {HeadDim::k64, 14, 2, 0, 16, 4096, 8192, 128, 1}, + {HeadDim::k64, 14, 14, 0, 16, 1, 0, 16, 7, 0}, + {HeadDim::k64, 14, 14, 0, 16, 16, 0, 16, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 64, 0, 4096, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 0, 4096, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 256, 512, 4096, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 512, 512, 14, 0}, + {HeadDim::k64, 14, 14, 0, 16, 1024, 2048, 128, 2, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 2048, 4096, 2, 0}, + {HeadDim::k64, 14, 14, 0, 16, 4096, 8192, 4096, 2, 0}, {HeadDim::k64, 16, 2, 0, 16, 1, 0, 16, 2}, {HeadDim::k64, 16, 2, 0, 16, 1, 512, 32, 2}, {HeadDim::k64, 16, 2, 0, 16, 1, 2048, 16, 2}, @@ -1672,6 +2141,33 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k64, 16, 2, 0, 16, 1024, 8192, 128, 2}, {HeadDim::k64, 16, 2, 0, 16, 4096, 2048, 512, 2}, {HeadDim::k64, 16, 2, 0, 16, 4096, 8192, 512, 1}, + {HeadDim::k64, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 64, 0, 32, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 0, 3072, 16, 0}, + {HeadDim::k64, 16, 16, 0, 16, 256, 512, 3072, 4, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 512, 3072, 1, 0}, + {HeadDim::k64, 16, 16, 0, 16, 1024, 2048, 3072, 1, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 2048, 3072, 2, 0}, + {HeadDim::k64, 16, 16, 0, 16, 4096, 8192, 3072, 2, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1, 0, 16, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 16, 0, 64, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 64, 0, 64, 20, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 0, 3072, 4, 0}, + {HeadDim::k64, 20, 20, 0, 16, 256, 512, 3072, 1, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 512, 3072, 1, 0}, + {HeadDim::k64, 20, 20, 0, 16, 1024, 2048, 3072, 2, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 2048, 3072, 2, 0}, + {HeadDim::k64, 20, 20, 0, 16, 4096, 8192, 3072, 2, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1, 0, 16, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 16, 0, 128, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 64, 0, 512, 24, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 0, 128, 8, 0}, + {HeadDim::k64, 24, 24, 0, 16, 256, 512, 3072, 4, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 512, 512, 12, 0}, + {HeadDim::k64, 24, 24, 0, 16, 1024, 2048, 3072, 4, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 2048, 3072, 2, 0}, + {HeadDim::k64, 24, 24, 0, 16, 4096, 8192, 3072, 2, 0}, {HeadDim::k64, 32, 8, 0, 16, 1, 0, 16, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 512, 16, 8}, {HeadDim::k64, 32, 8, 0, 16, 1, 2048, 16, 8}, @@ -1714,6 +2210,33 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k64, 32, 32, 0, 16, 1024, 8192, 512, 1}, {HeadDim::k64, 32, 32, 0, 16, 4096, 2048, 512, 4}, {HeadDim::k64, 32, 32, 0, 16, 4096, 8192, 256, 2}, + {HeadDim::k64, 64, 64, 0, 16, 1, 0, 32, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 16, 0, 32, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 64, 0, 32, 64, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 0, 1536, 8, 0}, + {HeadDim::k64, 64, 64, 0, 16, 256, 512, 1536, 32, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 512, 512, 4, 0}, + {HeadDim::k64, 64, 64, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k64, 64, 64, 0, 16, 4096, 8192, 1536, 4, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 16, 0, 64, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 64, 0, 1024, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 0, 1024, 16, 0}, + {HeadDim::k72, 16, 16, 0, 16, 256, 512, 3072, 1, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 512, 3072, 1, 0}, + {HeadDim::k72, 16, 16, 0, 16, 1024, 2048, 3072, 1, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 2048, 3072, 2, 0}, + {HeadDim::k72, 16, 16, 0, 16, 4096, 8192, 3072, 1, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1, 0, 32, 8, 0}, + {HeadDim::k80, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 64, 0, 16, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 0, 16, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 256, 512, 1536, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 512, 1536, 16, 0}, + {HeadDim::k80, 16, 16, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 2048, 1536, 4, 0}, + {HeadDim::k80, 16, 16, 0, 16, 4096, 8192, 1536, 2, 0}, {HeadDim::k96, 32, 32, 0, 16, 1, 0, 64, 2}, {HeadDim::k96, 32, 32, 0, 16, 1, 512, 128, 32}, {HeadDim::k96, 32, 32, 0, 16, 1, 2048, 64, 32}, @@ -1735,6 +2258,24 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k96, 32, 32, 0, 16, 1024, 8192, 2048, 2}, {HeadDim::k96, 32, 32, 0, 16, 4096, 2048, 2048, 4}, {HeadDim::k96, 32, 32, 0, 16, 4096, 8192, 2048, 1}, + {HeadDim::k128, 4, 4, 0, 16, 1, 0, 64, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 16, 0, 128, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 64, 0, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 0, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 256, 512, 768, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 512, 1536, 4, 0}, + {HeadDim::k128, 4, 4, 0, 16, 1024, 2048, 1536, 1, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 2048, 1536, 1, 0}, + {HeadDim::k128, 4, 4, 0, 16, 4096, 8192, 1536, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1, 0, 64, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 16, 0, 768, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 64, 0, 256, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 0, 768, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 256, 512, 3072, 8, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 512, 3072, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 1024, 2048, 3072, 1, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 2048, 3072, 4, 0}, + {HeadDim::k128, 8, 8, 0, 16, 4096, 8192, 3072, 1, 0}, {HeadDim::k128, 10, 10, 0, 16, 1, 0, 64, 10}, {HeadDim::k128, 10, 10, 0, 16, 1, 512, 64, 10}, {HeadDim::k128, 10, 10, 0, 16, 1, 2048, 16, 2}, @@ -1756,6 +2297,15 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k128, 10, 10, 0, 16, 1024, 8192, 256, 2}, {HeadDim::k128, 10, 10, 0, 16, 4096, 2048, 256, 10}, {HeadDim::k128, 10, 10, 0, 16, 4096, 8192, 256, 1}, + {HeadDim::k128, 12, 12, 0, 16, 1, 0, 16, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 16, 0, 768, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 64, 0, 1024, 12, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 0, 512, 3, 0}, + {HeadDim::k128, 12, 12, 0, 16, 256, 512, 512, 3, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 512, 1536, 1, 0}, + {HeadDim::k128, 12, 12, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 2048, 1536, 3, 0}, + {HeadDim::k128, 12, 12, 0, 16, 4096, 8192, 1536, 2, 0}, {HeadDim::k128, 16, 2, 0, 16, 1, 0, 32, 2}, {HeadDim::k128, 16, 2, 0, 16, 1, 512, 16, 2}, {HeadDim::k128, 16, 2, 0, 16, 1, 2048, 16, 2}, @@ -1798,6 +2348,15 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k128, 16, 8, 0, 16, 1024, 8192, 256, 1}, {HeadDim::k128, 16, 8, 0, 16, 4096, 2048, 256, 2}, {HeadDim::k128, 16, 8, 0, 16, 4096, 8192, 256, 1}, + {HeadDim::k128, 16, 16, 0, 16, 1, 0, 16, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 16, 0, 32, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 64, 0, 1536, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 0, 32, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 256, 512, 1536, 16, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 512, 1536, 8, 0}, + {HeadDim::k128, 16, 16, 0, 16, 1024, 2048, 1536, 4, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k128, 16, 16, 0, 16, 4096, 8192, 1536, 2, 0}, {HeadDim::k128, 24, 8, 0, 16, 1, 0, 16, 4}, {HeadDim::k128, 24, 8, 0, 16, 1, 512, 16, 8}, {HeadDim::k128, 24, 8, 0, 16, 1, 2048, 16, 8}, @@ -1819,6 +2378,15 @@ inline constexpr PrefillBatchConfig kGfx1201PrefillBatchConfigs[] = { {HeadDim::k128, 24, 8, 0, 16, 1024, 8192, 256, 1}, {HeadDim::k128, 24, 8, 0, 16, 4096, 2048, 256, 2}, {HeadDim::k128, 24, 8, 0, 16, 4096, 8192, 256, 1}, + {HeadDim::k128, 24, 24, 0, 16, 1, 0, 16, 24, 0}, + {HeadDim::k128, 24, 24, 0, 16, 16, 0, 16, 24, 0}, + {HeadDim::k128, 24, 24, 0, 16, 64, 0, 256, 24, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 0, 1536, 4, 0}, + {HeadDim::k128, 24, 24, 0, 16, 256, 512, 1536, 12, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 512, 16, 1, 0}, + {HeadDim::k128, 24, 24, 0, 16, 1024, 2048, 1536, 2, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 2048, 1536, 2, 0}, + {HeadDim::k128, 24, 24, 0, 16, 4096, 8192, 1536, 1, 0}, {HeadDim::k128, 28, 4, 0, 16, 1, 0, 16, 4}, {HeadDim::k128, 28, 4, 0, 16, 1, 512, 16, 4}, {HeadDim::k128, 28, 4, 0, 16, 1, 2048, 16, 4}, diff --git a/csrc/rocm/rdna_prefill_configs.h b/csrc/rocm/rdna_prefill_configs.h index 8268efa25a07..49154b422f6b 100644 --- a/csrc/rocm/rdna_prefill_configs.h +++ b/csrc/rocm/rdna_prefill_configs.h @@ -28,6 +28,11 @@ struct PrefillConfig { int key_waves; // waves sharing a row tile, merged at the end int max_segments; // most KV segments per tile int head_group; // kv heads whose tiles interleave in the grid + // 0: every query token sees every key of its sequence (no window, no + // multimodal ranges); a build of its own. Entries leave it out when 1. + int causal = 1; + // 16-row tiles per wave, 1 or (non-causal) 2; left out when 1. + int row_tiles = 1; }; // One table per Arch (rdna_attn_types.h), each complete; a GPU family's @@ -36,6 +41,136 @@ struct PrefillConfig { // line in PrefillConfig order. // clang-format off inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 1, 2, 4, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 0, 1, 2, 4, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 0, 1, 2, 1, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 8192, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 1, 2, 4, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 0, 1, 2, 4, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 0, 1, 2, 1, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 8192, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 0, 0, 2, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 4096, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 0, 0, 2, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 4096, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 1536, 4, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 1536, 4, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 4096, 4, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 4096, 4, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 0, 2, 4, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 0, 0, 2, 16, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1024, 16, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 4096, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 0, 2, 4, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 0, 0, 2, 16, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1024, 16, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 4096, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 2, 16, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 2, 16, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 0, 0, 2, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 1536, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 4096, 4, 16, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 0, 0, 2, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 1536, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 4096, 4, 16, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 1024, 4, 16, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 1024, 4, 16, 1, 1, 0, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 16, 8}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 0, 0, 2, 4, 1}, @@ -52,6 +187,22 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 0, 1, 1, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 256, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 1536, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 1, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 256, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 1536, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 0, 1, 4, 1, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 16, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 4}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 128, 8, 64, 1, 1, 0, 1, 16, 4}, @@ -66,6 +217,40 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 1, 16, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 2048, 4, 32, 1, 1, 0, 1, 1, 4}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 2, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 768, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 2, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 0, 1, 4, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 768, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 8192, 4, 16, 1, 1, 0, 1, 16, 1, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 8192, 4, 16, 1, 1, 0, 1, 16, 1, 0, 2}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 4, 1}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 4}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 0, 0, 1, 4, 1}, @@ -88,6 +273,56 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 1, 8}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 256, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 1536, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 8192, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 256, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 1536, 4, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 8192, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 0, 1, 4, 2, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 0, 1, 4, 2, 0, 2}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 64, 1, 0, 0, 2, 1, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 2, 1, 2}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 1, 2}, @@ -102,6 +337,44 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 0, 0, 1, 4, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 0, 0, 1, 1, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 16, 4}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 2, 4, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 2, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 512, 16, 32, 1, 1, 0, 2, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 4096, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 2, 4, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 2, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 512, 16, 32, 1, 1, 0, 2, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 4096, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 512, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 1024, 16, 16, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 1024, 16, 16, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 4, 4}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 0, 1, 4, 8}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 0, 0, 1, 4, 1}, @@ -112,6 +385,26 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 0, 0, 1, 4, 1}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 1, 0, 2, 1, 4}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 0, 2, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 1536, 16, 32, 1, 1, 1, 2, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 4096, 8, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 2, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 1536, 16, 32, 1, 1, 1, 2, 4, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 4096, 8, 32, 1, 1, 1, 1, 1, 4, 0}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 1, 0, 2, 16, 2}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 0, 1, 16, 1}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 16, 8}, @@ -136,6 +429,26 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 1, 16, 1}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 512, 4, 16, 1, 1, 0, 1, 1, 8}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 8}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 128, 8, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 512, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1024, 16, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 8192, 16, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 128, 8, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 512, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1024, 16, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1536, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 8192, 16, 32, 1, 1, 0, 1, 16, 8, 0}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 32, 8, 16, 1, 1, 0, 1, 4, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 8}, @@ -144,6 +457,20 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 32, 8, 16, 1, 1, 0, 1, 4, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 8}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 16, 8}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 1536, 16, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 8192, 16, 32, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 1536, 16, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 8192, 16, 32, 1, 1, 0, 1, 1, 1, 0}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 4}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 4}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 2}, @@ -473,6 +800,164 @@ inline constexpr PrefillConfig kGfx1151PrefillConfigs[] = { }; inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 16, 16, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 1, 1, 16, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 0, 0, 2, 1, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 2, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 3072, 4, 16, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 4096, 8, 32, 1, 1, 0, 1, 16, 1, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 16, 16, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 1, 1, 16, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 0, 0, 2, 1, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 2, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 3072, 4, 16, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 4096, 8, 32, 1, 1, 0, 1, 16, 1, 0, 2}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 128, 8, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 0, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 0, 1, 16, 2, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 8192, 16, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 0, 1, 16, 2, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 8192, 16, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 2, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 4096, 8, 16, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 8192, 16, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 2, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 4096, 8, 16, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 8192, 16, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 2, 1, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 2, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 256, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 1024, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 2048, 8, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 8192, 16, 64, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 2, 1, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 2, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 256, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 1024, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 2048, 8, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 8192, 16, 64, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 1, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 1, 1, 1, 1, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 1, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 1, 1, 1, 1, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 16, 8, 16, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 0, 1, 2, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1024, 4, 64, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 2, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 1, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 16, 8, 16, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 0, 1, 2, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 2, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 1, 1, 4, 1, 0, 2}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 64, 8, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 64, 8, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 16, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 256, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 2048, 8, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 16, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 256, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 2048, 8, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 4096, 4, 32, 1, 1, 1, 1, 1, 4, 0, 2}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 16, 4, 64, 1, 0, 0, 2, 4, 8}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 1, 0, 1, 16, 1}, @@ -485,6 +970,20 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 1, 0, 1, 1, 2}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 0, 0, 1, 1, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 1, 1, 1}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 4096, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 1, 0, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 4096, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 64, 1, 0, 0, 2, 4, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 4}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 1, 0, 1, 16, 1}, @@ -497,6 +996,50 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 1, 1, 1}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 2048, 4, 32, 1, 0, 0, 1, 16, 1}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 2048, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 8192, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 2, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 1536, 4, 32, 1, 1, 1, 1, 16, 4, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 2048, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 8192, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 2, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 16, 2, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 16, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 1, 1, 1, 8, 0, 2}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 1, 1, 1, 8, 0, 2}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 16, 8, 64, 1, 0, 0, 2, 1, 8}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 8}, @@ -517,6 +1060,58 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 0, 0, 1, 16, 1}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 1}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 4, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 1, 1, 4, 8, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 512, 8, 32, 1, 1, 0, 1, 1, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1536, 16, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 2048, 16, 32, 1, 1, 0, 1, 4, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 512, 8, 32, 1, 1, 0, 1, 1, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 32, 1, 1, 0, 1, 1, 1, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1536, 16, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 2048, 16, 32, 1, 1, 0, 1, 4, 2, 0, 2}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 0, 1, 1, 8, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 2, 1, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 32, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 8192, 16, 32, 1, 1, 0, 1, 16, 8, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 2, 1, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 4, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 32, 1, 1, 0, 1, 16, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 32, 1, 1, 0, 1, 1, 4, 0, 2}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 8192, 16, 32, 1, 1, 0, 1, 16, 8, 0, 2}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 1, 2}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 1, 4}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 0, 0, 1, 16, 1}, @@ -525,6 +1120,32 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 1, 1, 4}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 0, 0, 1, 16, 1}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 8192, 8, 64, 1, 0, 0, 1, 16, 2}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 64, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 128, 8, 64, 1, 0, 0, 2, 1, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 768, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 64, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 0, 0, 2, 1, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 768, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 4, 4, 0}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 1, 8}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 256, 8, 64, 1, 1, 0, 2, 1, 2}, @@ -533,6 +1154,18 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 2, 1, 2}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 4}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 1, 1, 4, 0}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 4, 8}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 32, 8, 32, 1, 1, 0, 1, 4, 2}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 1, 0, 2, 1, 4}, @@ -553,6 +1186,24 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 1, 16, 2}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 1, 16, 1}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 2048, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 3072, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 3072, 4, 32, 1, 1, 0, 1, 16, 4, 0}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 4}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8}, @@ -563,6 +1214,14 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 1, 8}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 1, 2}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 1}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 8192, 16, 32, 1, 1, 1, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 8192, 16, 32, 1, 1, 1, 1, 1, 4, 0}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 16, 2}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 0, 1, 16, 4}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 4}, @@ -810,6 +1469,150 @@ inline constexpr PrefillConfig kGfx1150PrefillConfigs[] = { }; inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 256, 8, 64, 1, 1, 0, 2, 16, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1024, 4, 64, 1, 1, 1, 2, 4, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 3072, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 4096, 4, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kFp16, 16, 8192, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 2, 16, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 1, 2, 4, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 3072, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 4096, 4, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k32, 4, 4, 0, Dtype::kBf16, 16, 8192, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 768, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kFp16, 16, 4096, 8, 64, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 768, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k32, 8, 8, 0, Dtype::kBf16, 16, 4096, 8, 64, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kFp16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 4, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 16, 16, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kFp16, 16, 1024, 8, 64, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 4, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k32, 32, 32, 0, Dtype::kBf16, 16, 1024, 8, 64, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 512, 8, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kFp16, 16, 1536, 16, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 512, 8, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k48, 32, 32, 0, Dtype::kBf16, 16, 1536, 16, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 256, 4, 64, 1, 0, 0, 2, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1024, 4, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 4096, 16, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kFp16, 16, 8192, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 256, 4, 64, 1, 0, 0, 2, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 4096, 16, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 4, 4, 0, Dtype::kBf16, 16, 8192, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 1, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 256, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1024, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kFp16, 16, 8192, 8, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 1, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 256, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 2, 4, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 1, 1, 16, 2, 0}, + {HeadDim::k64, 6, 6, 0, Dtype::kBf16, 16, 8192, 8, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 1024, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 1536, 8, 32, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 2048, 16, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 3072, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kFp16, 16, 8192, 8, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 1536, 8, 32, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 2048, 16, 32, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 3072, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 8, 8, 0, Dtype::kBf16, 16, 8192, 8, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 128, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 2048, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kFp16, 16, 4096, 8, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 128, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 1, 1, 4, 4, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 2048, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k64, 12, 12, 0, Dtype::kBf16, 16, 4096, 8, 32, 1, 1, 0, 1, 16, 4, 0}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 16, 8}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kFp16, 16, 64, 8, 64, 1, 0, 0, 2, 4, 1}, @@ -826,6 +1629,16 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 1024, 4, 64, 1, 1, 0, 1, 1, 1}, {HeadDim::k64, 14, 2, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 1, 2}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 1536, 4, 64, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kFp16, 16, 4096, 8, 32, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 1536, 4, 64, 1, 1, 0, 1, 4, 1, 0}, + {HeadDim::k64, 14, 14, 0, Dtype::kBf16, 16, 4096, 8, 32, 1, 1, 0, 1, 4, 2, 0}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 16, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 4}, {HeadDim::k64, 16, 2, 0, Dtype::kFp16, 16, 128, 8, 64, 1, 1, 0, 1, 16, 4}, @@ -836,6 +1649,44 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 128, 8, 64, 1, 1, 0, 1, 16, 4}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 0, 0, 1, 4, 8}, {HeadDim::k64, 16, 2, 0, Dtype::kBf16, 16, 512, 8, 32, 1, 2, 1, 1, 4, 4}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 2, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 1536, 4, 64, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 2048, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 2, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 1536, 4, 64, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k64, 16, 16, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 1536, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 2048, 16, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 1536, 4, 64, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 2048, 16, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 20, 20, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 2, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 768, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k64, 24, 24, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 0, 1, 16, 4, 0}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 4, 1}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 32, 8, 64, 1, 0, 0, 2, 16, 4}, {HeadDim::k64, 32, 8, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 0, 0, 1, 4, 1}, @@ -856,6 +1707,44 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 0, 0, 1, 16, 8}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 2, 1, 1, 4, 4}, {HeadDim::k64, 32, 32, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 128, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kFp16, 16, 1536, 8, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 0, 0, 1, 4, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k64, 64, 64, 0, Dtype::kBf16, 16, 1536, 8, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 256, 8, 64, 1, 1, 0, 2, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 768, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 1536, 16, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 256, 8, 64, 1, 1, 0, 2, 1, 2, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 768, 8, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1024, 8, 16, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 1536, 16, 64, 1, 1, 0, 1, 4, 4, 0}, + {HeadDim::k72, 16, 16, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kFp16, 16, 1536, 16, 64, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 64, 4, 64, 1, 1, 0, 1, 1, 1, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 512, 4, 64, 1, 1, 0, 1, 4, 8, 0}, + {HeadDim::k80, 16, 16, 0, Dtype::kBf16, 16, 1536, 16, 64, 1, 1, 1, 1, 1, 8, 0}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 16, 4, 64, 1, 0, 0, 2, 1, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 2, 1, 2}, {HeadDim::k96, 32, 32, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 1, 1, 2}, @@ -870,6 +1759,40 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 0, 0, 1, 4, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 0, 0, 1, 1, 8}, {HeadDim::k96, 32, 32, 0, Dtype::kBf16, 16, 2048, 8, 64, 1, 1, 0, 1, 16, 4}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 0, 0, 2, 1, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 2, 16, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 768, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 3072, 8, 32, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kFp16, 16, 4096, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 0, 0, 2, 1, 1, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 2, 16, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 768, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 3072, 8, 32, 1, 1, 0, 2, 1, 8, 0}, + {HeadDim::k128, 4, 4, 0, Dtype::kBf16, 16, 4096, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 1536, 16, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kFp16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 64, 4, 16, 1, 0, 0, 1, 16, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 1536, 16, 64, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k128, 8, 8, 0, Dtype::kBf16, 16, 3072, 8, 64, 1, 1, 0, 1, 1, 8, 0}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 0, 0, 2, 4, 4}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 64, 4, 16, 1, 1, 0, 1, 4, 8}, {HeadDim::k128, 10, 10, 0, Dtype::kFp16, 16, 128, 8, 32, 1, 0, 0, 1, 4, 1}, @@ -880,6 +1803,20 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 128, 8, 32, 1, 0, 0, 1, 4, 1}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 1, 1, 1, 4, 4}, {HeadDim::k128, 10, 10, 0, Dtype::kBf16, 16, 512, 8, 64, 1, 1, 0, 1, 1, 8}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 4, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 32, 4, 32, 1, 0, 0, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 64, 4, 32, 1, 0, 0, 1, 1, 1, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 1, 1, 1, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 12, 12, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 4, 0}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 16, 4, 32, 1, 1, 0, 2, 16, 2}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 0, 1, 16, 1}, {HeadDim::k128, 16, 2, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 16, 8}, @@ -902,6 +1839,20 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 128, 4, 16, 1, 1, 0, 1, 16, 8}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 1, 1, 1, 4, 4}, {HeadDim::k128, 16, 8, 0, Dtype::kBf16, 16, 512, 8, 32, 1, 1, 1, 1, 4, 4}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 32, 4, 16, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 768, 8, 32, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 0, 1, 1, 1, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 32, 4, 16, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 1, 1, 1, 2, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 512, 4, 32, 1, 1, 0, 1, 1, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 768, 8, 32, 1, 1, 0, 2, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1024, 4, 32, 1, 1, 0, 1, 16, 8, 0}, + {HeadDim::k128, 16, 16, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 16, 8, 0}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 4, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 32, 8, 16, 1, 1, 0, 1, 4, 1}, {HeadDim::k128, 24, 8, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 8}, @@ -914,6 +1865,16 @@ inline constexpr PrefillConfig kGfx1201PrefillConfigs[] = { {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 8}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 256, 8, 32, 1, 1, 1, 1, 4, 4}, {HeadDim::k128, 24, 8, 0, Dtype::kBf16, 16, 512, 8, 32, 1, 1, 1, 1, 4, 4}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kFp16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 16, 4, 16, 1, 1, 1, 1, 1, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 256, 4, 32, 1, 1, 0, 1, 16, 1, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 768, 4, 16, 1, 1, 1, 1, 16, 4, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 1024, 4, 16, 1, 1, 1, 1, 4, 8, 0}, + {HeadDim::k128, 24, 24, 0, Dtype::kBf16, 16, 1536, 8, 64, 1, 1, 0, 1, 1, 4, 0}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 16, 4, 16, 1, 1, 0, 1, 16, 4}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 64, 4, 32, 1, 1, 0, 1, 4, 4}, {HeadDim::k128, 28, 4, 0, Dtype::kFp16, 16, 128, 4, 32, 1, 1, 0, 1, 4, 2}, diff --git a/csrc/rocm/rdna_registry.cu b/csrc/rocm/rdna_registry.cu index 85776f103753..556d18797b24 100644 --- a/csrc/rocm/rdna_registry.cu +++ b/csrc/rocm/rdna_registry.cu @@ -190,7 +190,7 @@ std::vector rdna_prefill_configs() { {static_cast(c.head_dim), c.num_q_heads, c.num_kv_heads, c.window, static_cast(c.dtype), c.page_size, c.min_query_len, c.waves, c.key_tile, c.head_dim_split, c.v_t, c.double_buffer, c.key_waves, - c.max_segments, c.head_group}); + c.max_segments, c.head_group, c.causal, c.row_tiles}); } return fields; } @@ -209,7 +209,7 @@ std::vector rdna_prefill_batch_configs() { return c.head_dim == b.head_dim && c.num_q_heads == b.num_q_heads && c.num_kv_heads == b.num_kv_heads && c.window == b.window && c.page_size == b.page_size && c.dtype == dtype && - c.min_query_len == b.min_query_len; + c.min_query_len == b.min_query_len && c.causal == b.causal; }); TORCH_CHECK(built, "prefill batch cell of ", b.num_q_heads, "/", b.num_kv_heads, "/", static_cast(b.head_dim), " window ", @@ -217,10 +217,10 @@ std::vector rdna_prefill_batch_configs() { " names no prefill build with min_query_len ", b.min_query_len, ", bf16 ", dtype == Dtype::kBf16); } - fields.insert( - fields.end(), - {static_cast(b.head_dim), b.num_q_heads, b.num_kv_heads, b.window, - b.page_size, b.min_mean_m, b.min_keys, b.min_query_len, b.head_group}); + fields.insert(fields.end(), + {static_cast(b.head_dim), b.num_q_heads, b.num_kv_heads, + b.window, b.page_size, b.min_mean_m, b.min_keys, + b.min_query_len, b.head_group, b.causal}); } return fields; } @@ -361,6 +361,7 @@ void rdna_prefill_attn(int64_t config, torch::Tensor& q, config >= 0 && config < t.num_prefill && t.prefill_launchers[config], "no prefill build ", config); const PrefillConfig& c = t.prefill[config]; + TORCH_CHECK(c.causal, "a non-causal build runs through rdna_encoder_attn"); TORCH_CHECK(q.is_contiguous() && out.is_contiguous(), "q and out must be contiguous"); TORCH_CHECK( @@ -429,7 +430,8 @@ void rdna_prefill_attn(int64_t config, torch::Tensor& q, // As PrefillVariant.scratch_rows. const int64_t target = 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; - const int64_t block_rows = c.waves / (c.head_dim_split * c.key_waves) * 16; + const int64_t block_rows = + c.waves / (c.head_dim_split * c.key_waves) * 16 * c.row_tiles; const int64_t rows = 2 * target * block_rows; TORCH_CHECK(partial_o.numel() >= rows * d && partial_ml.numel() >= 2 * rows && @@ -448,3 +450,100 @@ void rdna_prefill_attn(int64_t config, torch::Tensor& q, block_table.dim() == 2 ? (int)block_table.stride(0) : 0, (int)head_group}); } + +void rdna_encoder_attn(int64_t config, torch::Tensor& q, torch::Tensor& k, + torch::Tensor& v, torch::Tensor& out, + torch::Tensor& partial_o, torch::Tensor& partial_ml, + torch::Tensor& counters, torch::Tensor& seq_lens, + int64_t max_seq_len, double scale, + const std::optional& plan, + const std::optional& query_start_loc, + int64_t head_group) { + using namespace rdna; + const Table t = device_table(); + TORCH_CHECK( + config >= 0 && config < t.num_prefill && t.prefill_launchers[config], + "no prefill build ", config); + const PrefillConfig& c = t.prefill[config]; + TORCH_CHECK(!c.causal, "a causal build runs through rdna_prefill_attn"); + const int d = static_cast(c.head_dim); + // (tokens, heads, D), each token's heads contiguous; any token stride. + auto check = [&](const torch::Tensor& x, int heads, const char* name) { + TORCH_CHECK(x.dim() == 3 && x.size(0) == q.size(0) && x.size(1) == heads && + x.size(2) == d && x.stride(2) == 1 && x.stride(1) == d, + name, " must be (", q.size(0), ", ", heads, ", ", d, + ") with each token's heads contiguous, got ", x.sizes(), + " strides ", x.strides()); + }; + check(q, c.num_q_heads, "q"); + check(k, c.num_kv_heads, "k"); + check(v, c.num_kv_heads, "v"); + TORCH_CHECK(out.is_contiguous() && out.sizes() == q.sizes(), + "out must be contiguous and shaped as q"); + const auto dtype = c.dtype == Dtype::kBf16 ? at::kBFloat16 : at::kHalf; + TORCH_CHECK(q.scalar_type() == dtype && k.scalar_type() == dtype && + v.scalar_type() == dtype && out.scalar_type() == dtype, + "kernel built for ", dtype); + TORCH_CHECK(seq_lens.scalar_type() == torch::kInt32 && seq_lens.numel() >= 1, + "seq_lens must be int32, a length per sequence"); + TORCH_CHECK(plan.has_value() == query_start_loc.has_value(), + "a batch takes both a plan and query_start_loc"); + if (plan) { + TORCH_CHECK(plan->scalar_type() == torch::kInt32 && plan->dim() == 2 && + plan->size(1) == 4 && plan->is_contiguous(), + "plan must be plan_prefill's contiguous int32 [entries, 4]"); + TORCH_CHECK(query_start_loc->scalar_type() == torch::kInt32 && + query_start_loc->numel() >= 2 && + seq_lens.numel() >= query_start_loc->numel() - 1, + "query_start_loc must be int32, one entry per sequence and " + "one, with a length per sequence"); + TORCH_CHECK(head_group >= 0 && + (head_group == 0 || c.num_kv_heads % head_group == 0), + "head_group must divide the kv heads"); + } else { + TORCH_CHECK(q.size(0) > 0 && max_seq_len == q.size(0), + "one sequence: max_seq_len is its tokens"); + } + TORCH_CHECK(partial_o.scalar_type() == torch::kFloat32 && + partial_ml.scalar_type() == torch::kFloat32 && + counters.scalar_type() == torch::kInt32, + "scratch must be float32 partials and int32 counters"); + check_device(q, + {&k, &v, &out, &partial_o, &partial_ml, &counters, &seq_lens}); + if (plan) check_device(q, {&*plan, &*query_start_loc}); + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + if (c.max_segments > 1) { + // As PrefillVariant.scratch_rows. + const int64_t target = + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int64_t block_rows = + c.waves / (c.head_dim_split * c.key_waves) * 16 * c.row_tiles; + const int64_t rows = 2 * target * block_rows; + TORCH_CHECK(partial_o.numel() >= rows * d && + partial_ml.numel() >= 2 * rows && + counters.numel() >= target, + "scratch too small for this build: needs ", rows * d, ", ", + 2 * rows, ", ", target, " elements"); + } + PrefillArgs args{}; + args.q = q.data_ptr(); + args.kv = k.data_ptr(); + args.out = out.data_ptr(); + args.seq_lens = seq_lens.data_ptr(); + args.partial_o = partial_o.data_ptr(); + args.partial_ml = partial_ml.data_ptr(); + args.counters = counters.data_ptr(); + args.num_tokens = (int)q.size(0); + args.max_seq_len = (int)max_seq_len; + args.scale = (float)scale; + args.stream = at::cuda::getCurrentCUDAStream(); + args.plan = plan ? plan->data_ptr() : nullptr; + args.query_start = plan ? query_start_loc->data_ptr() : nullptr; + args.num_entries = plan ? (int)plan->size(0) : 0; + args.head_group = (int)head_group; + args.v = v.data_ptr(); + args.q_stride = (int)q.stride(0); + args.k_stride = (int)k.stride(0); + args.v_stride = (int)v.stride(0); + t.prefill_launchers[config](args); +} diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index f765de43e92a..73b04a6c29b2 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -212,6 +212,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor? plan=None, Tensor? query_start_loc=None, int head_group=0) -> " "()"); rocm_ops.impl("rdna_prefill_attn", torch::kCUDA, &rdna_prefill_attn); + // The non-causal builds, on an encoder's Q, K and V in place. + rocm_ops.def( + "rdna_encoder_attn(int config, Tensor q, Tensor k, Tensor v, " + "Tensor! out, Tensor! partial_o, Tensor! partial_ml, Tensor! counters, " + "Tensor seq_lens, int max_seq_len, float scale, Tensor? plan=None, " + "Tensor? query_start_loc=None, int head_group=0) -> ()"); + rocm_ops.impl("rdna_encoder_attn", torch::kCUDA, &rdna_encoder_attn); #endif // Custom attention op diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index b36f5776da25..66b739b48ebf 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -49,6 +49,21 @@ def test_mm_encoder_attn_backend_invalid(): MultiModalConfig(mm_encoder_attn_backend="not_a_backend") # type: ignore[arg-type] +def test_rocm_rdna_mm_encoder_rejects_encoder_cudagraph(): + """ROCM_RDNA_ATTN plans each launch on the host, so a captured encoder + graph would replay stale plans: the combination fails at startup.""" + from vllm.config.compilation import CompilationConfig + + model_config = ModelConfig( + _MULTIMODAL_MODEL, mm_encoder_attn_backend="ROCM_RDNA_ATTN" + ) + with pytest.raises(ValueError, match="ROCM_RDNA_ATTN is incompatible"): + VllmConfig( + model_config=model_config, + compilation_config=CompilationConfig(cudagraph_mm_encoder=True), + ) + + def test_mm_hasher_algorithm_invalid(): with pytest.raises(ValueError, match="mm_hasher_algorithm"): MultiModalConfig(mm_hasher_algorithm="md5") # type: ignore[arg-type] diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index 5a9c8c68d1e6..38853000b3e5 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -488,3 +488,161 @@ def test_mha_attn_fp8_rejects_wrong_backend(default_vllm_config): MMEncoderAttention(4, 64) finally: vllm_config.model_config = old_model_config + + +def _rdna_mm_encoder(head_size: int, monkeypatch): + """An MMEncoderAttention on ROCM_RDNA_ATTN, and the calls its kernels + served (not TRITON_ATTN's fallback).""" + from vllm.platforms.rocm import on_gfx115x, on_rdna4 + from vllm.v1.attention.ops import rocm_rdna_encoder + + if not current_platform.is_rocm() or not (on_gfx115x() or on_rdna4()): + pytest.skip("ROCM_RDNA_ATTN's encoder kernels: gfx115x, RDNA4") + served = [] + forward = rocm_rdna_encoder.forward + + def spy(q, k, v, seg_lens, scale): + served.append(list(seg_lens)) + return forward(q, k, v, seg_lens, scale) + + monkeypatch.setattr(rocm_rdna_encoder, "forward", spy) + with patch( + "vllm.model_executor.layers.attention.mm_encoder_attention." + "get_vit_attn_backend", + return_value=AttentionBackendEnum.ROCM_RDNA_ATTN, + ): + attn = MMEncoderAttention(16, head_size, scale=head_size**-0.5) + assert attn.attn_backend == AttentionBackendEnum.ROCM_RDNA_ATTN + return attn, served + + +@pytest.mark.parametrize( + "segments", + [[4096], [577] * 5, [1024, 257, 9, 4, 7, 64, 8, 1], [4, 8, 4, 1]], + ids=["one", "equal", "mixed", "short"], +) +@pytest.mark.parametrize("head_size", [72, 64, 80]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_mha_attn_rdna_varlen_forward( + default_vllm_config, segments, head_size, dtype, monkeypatch +): + """ROCM_RDNA_ATTN serves an encoder's segments, however short, in one + non-causal prefill launch, as the reference.""" + set_random_seed(0) + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + attn, served = _rdna_mm_encoder(head_size, monkeypatch) + q, k, v = (torch.randn(1, sum(segments), 16, head_size) for _ in range(3)) + cu_seqlens = torch.tensor([0, *itertools.accumulate(segments)], dtype=torch.int32) + out = attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=torch.tensor(max(segments))) + assert served == [segments], "the call fell back to TRITON_ATTN" + ref = torch.cat( + [ + ref_attention(q_i.float(), k_i.float(), v_i.float(), head_size**-0.5) + for q_i, k_i, v_i in zip( + *(torch.split(x, segments, dim=1) for x in (q, k, v)) + ) + ], + dim=1, + ) + torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize("head_size", [72, 64]) +def test_mha_attn_rdna_forward_without_cu_seqlens( + default_vllm_config, head_size, monkeypatch +): + """Without cu_seqlens each batch row is a segment.""" + set_random_seed(0) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.float16) + attn, served = _rdna_mm_encoder(head_size, monkeypatch) + q, k, v = (torch.randn(2, 320, 16, head_size) for _ in range(3)) + out = attn(q, k, v) + assert served == [[320, 320]] + ref = ref_attention(q.float(), k.float(), v.float(), scale=head_size**-0.5) + torch.testing.assert_close(out.float().view_as(ref), ref, atol=1e-2, rtol=1e-2) + + +def test_mha_attn_rdna_falls_back_to_triton(default_vllm_config, monkeypatch): + """A head size without builds runs on TRITON_ATTN, with the same result.""" + set_random_seed(0) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.float16) + attn, served = _rdna_mm_encoder(40, monkeypatch) + q, k, v = (torch.randn(1, 300, 16, 40) for _ in range(3)) + out = attn(q, k, v) + assert served == [] + ref = ref_attention(q.float(), k.float(), v.float(), scale=40**-0.5) + torch.testing.assert_close(out.float().view_as(ref), ref, atol=1e-2, rtol=1e-2) + + +def test_rdna_segment_lengths_kept_per_live_tensor(): + """ROCM_RDNA_ATTN reads a cu_seqlens tensor's lengths once, whether noted + from the host or read back, and forgets them only with that tensor.""" + from vllm.v1.attention.ops import vit_attn_wrappers as W + + noted = torch.tensor([0, 3, 7], dtype=torch.int32) + W.remember_rdna_cu_seqlens(noted, np.array([0, 3, 7])) + read = [torch.tensor([0, i + 1], dtype=torch.int32) for i in range(16)] + for t in read: + assert W._rdna_segment_lengths(t) is W._rdna_segment_lengths(t) + assert W._rdna_segment_lengths(noted) == [3, 4] + gone = id(read.pop(0)) + assert gone not in W._RDNA_SEG_LENS + assert id(noted) in W._RDNA_SEG_LENS + + +def _varlen_reference(q, k, v, segments, scale): + """Attention within each segment, in float.""" + return torch.cat( + [ + ref_attention(q_i.float(), k_i.float(), v_i.float(), scale) + for q_i, k_i, v_i in zip( + *(torch.split(x, segments, dim=1) for x in (q, k, v)) + ) + ], + dim=1, + ) + + +def _check_cuda_graph_replay(attn, head_size, capture, replays): + """Capture attn on cu_seqlens of the `capture` segments, with their total + as max_seqlen (a bound, as encoder CUDA graphs pass), then replay it on + each list of `replays` (the same total, padded with empty segments).""" + tokens, slots = sum(capture), max(map(len, [capture, *replays])) + q, k, v = (torch.randn(1, tokens, 16, head_size) for _ in range(3)) + + def bounds(segments): + padded = [*segments, *[0] * (slots - len(segments))] + return torch.tensor([0, *itertools.accumulate(padded)], dtype=torch.int32) + + cu_seqlens = bounds(capture) + max_seqlen = torch.tensor(tokens, dtype=torch.int32, device="cpu") + attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + for segments in replays: + cu_seqlens.copy_(bounds(segments)) + graph.replay() + ref = _varlen_reference(q, k, v, segments, head_size**-0.5) + torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA graphs") +def test_mha_attn_triton_varlen_cuda_graph(default_vllm_config): + """TRITON_ATTN's varlen path is captured once and replayed on other + segment lists: nothing it reads on the host is baked into the graph.""" + set_random_seed(0) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.float16) + with patch( + "vllm.model_executor.layers.attention.mm_encoder_attention." + "get_vit_attn_backend", + return_value=AttentionBackendEnum.TRITON_ATTN, + ): + attn = MMEncoderAttention(16, 64, scale=64**-0.5) + _check_cuda_graph_replay( + attn, 64, [100, 150, 50], [[40, 220, 40], [300], [1, 2, 297]] + ) diff --git a/tests/kernels/attention/test_rocm_rdna_hip.py b/tests/kernels/attention/test_rocm_rdna_hip.py index f7b3c18e631f..93f5b8385fa1 100644 --- a/tests/kernels/attention/test_rocm_rdna_hip.py +++ b/tests/kernels/attention/test_rocm_rdna_hip.py @@ -28,11 +28,17 @@ @pytest.mark.parametrize( - "kind, shape_fields", - [("decode", 7), ("decode_batch", 7), ("prefill", 6), ("prefill_batch", 5)], + "kind, shape_fields, causal_field", + [ + ("decode", 7, None), + ("decode_batch", 7, None), + ("prefill", 6, 15), + ("prefill_batch", 5, 9), + ], ) -def test_gfx1150_configs_cover_the_same_shapes(kind, shape_fields): - """gfx1150's tables serve the shapes gfx1151's serve.""" +def test_gfx1150_configs_cover_the_same_shapes(kind, shape_fields, causal_field): + """gfx1150's tables serve the shapes gfx1151's serve, but for the + non-causal (mm-encoder) builds, gfx1151's only.""" csrc = Path(__file__).resolve().parents[3] / "csrc" / "rocm" text = (csrc / f"rdna_{kind}_configs.h").read_text() shapes = {} @@ -43,6 +49,12 @@ def test_gfx1150_configs_cover_the_same_shapes(kind, shape_fields): for m in map(_CONFIG_ENTRY.match, table.split("\n")) if m ] + if causal_field is not None: + entries = [ + e + for e in entries + if len(e) <= causal_field or e[causal_field] != "0" + ] shapes[name[1]] = {tuple(e[:shape_fields]) for e in entries} assert shapes["Gfx1151"] and shapes["Gfx1151"] == shapes["Gfx1150"] @@ -356,7 +368,7 @@ def test_unsupported_head_size_falls_back_to_triton(): impl = _impl() impl._rejected = None - impl.head_size, impl.num_heads, impl.num_kv_heads = 80, HQ, HKV + impl.head_size, impl.num_heads, impl.num_kv_heads = 160, HQ, HKV fits = impl._prepare( kv_cache=torch.empty(0, dtype=torch.float16), @@ -372,7 +384,7 @@ def test_unsupported_head_size_falls_back_to_triton(): max_seqlen_q=0, ) assert not fits - assert "head_size 80" in impl._rejected + assert "head_size 160" in impl._rejected def test_strided_query_falls_back_to_triton(): @@ -478,16 +490,18 @@ def test_rdna_valid_for_any_batch(): # ---------------------------------------------------------------- prefill -def _prefill_reference(q, kv, block_table, seq_len, window=0, mm_ranges=()): +def _prefill_reference( + q, kv, block_table, seq_len, window=0, mm_ranges=(), causal=True +): """Float reference of a prefill; a token in a multimodal range [start, end] - also sees the range's later keys.""" + also sees the range's later keys. Non-causal: every key.""" m, hq, hkv, hd = q.shape[0], q.shape[1], kv.shape[1], q.shape[2] pages = kv[block_table.long()] flat = pages.transpose(1, 2).reshape(-1, hkv, 2 * hd)[:seq_len].float() gqa = hq // hkv pos = torch.arange(seq_len, device=q.device).view(1, seq_len) lim = (seq_len - m + torch.arange(m, device=q.device)).view(m, 1) - masked = pos > lim + masked = pos > lim if causal else torch.zeros_like(pos > lim) for start, end in mm_ranges: q_in = (lim >= start) & (lim <= end) k_in = (pos >= start) & (pos <= end) @@ -575,8 +589,10 @@ def _prefill_inputs(v, m, seq_len, seed=0): def _prefill_unit(v): dt = "bf16" if v.dtype == torch.bfloat16 else "fp16" + mode = "" if v.causal else "_noncausal" return ( f"prefill_q{v.num_q_heads}_kv{v.num_kv_heads}_d{v.head_size}_w{v.window}_{dt}" + f"{mode}" ) @@ -584,6 +600,10 @@ def _prefill_rows(): return [v for rows in hip._prefill_table().values() for v in rows] +def _causal_prefill_rows(): + return [v for v in _prefill_rows() if v.causal] + + def test_prefill_configs_are_well_formed(): """No two prefill builds share a shape and step, every decode configuration has prefill builds, and the build of the largest step reached runs.""" @@ -591,7 +611,7 @@ def test_prefill_configs_are_well_formed(): pytest.skip("_rocm_C carries no RDNA variants") keys = { (v.num_q_heads, v.num_kv_heads, v.head_size, v.page_size, v.window) - + (v.dtype, v.min_query_len) + + (v.dtype, v.min_query_len, v.causal) for v in _prefill_rows() } assert len(keys) == len(_prefill_rows()), "two builds share a step" @@ -623,7 +643,9 @@ def test_every_listed_prefill_variant_is_built_in(): assert not missing, f"{len(missing)} prefill variants not in _rocm_C" -@pytest.mark.parametrize("unit", sorted({_prefill_unit(v) for v in _prefill_rows()})) +@pytest.mark.parametrize( + "unit", sorted({_prefill_unit(v) for v in _causal_prefill_rows()}) +) def test_built_in_prefill_variants_match_reference(unit): """Every built prefill variant matches the reference, with and without a prefix.""" @@ -668,7 +690,9 @@ def test_prefill_past_the_last_step(): v = next( rows[-1] for rows in hip._prefill_table().values() - if rows[-1].dtype == torch.float16 and rows[-1].head_size == 128 + if rows[-1].dtype == torch.float16 + and rows[-1].head_size == 128 + and rows[-1].causal ) module = hip.load_prefill(v) args = (v.num_q_heads, v.num_kv_heads, v.head_size, v.page_size, v.window) @@ -713,6 +737,210 @@ def test_prefill_window_never_reads_freed_pages(): _assert_prefill_close(out, q, *clean, seq_len, v.window, "freed pages") +# ---------------------------------------------------------------- non-causal + + +def _non_causal_rows(): + return [v for v in _prefill_rows() if not v.causal] + + +def _non_causal_inputs(v, lens, seed=0): + """Q, K and V of sequences of `lens` tokens, one after the other, as views + of one (tokens, 3, H, D) projection (a token stride of 3 H D).""" + torch.manual_seed(seed) + n = sum(lens) + qkv = torch.randn(n, 3, v.num_q_heads, v.head_size, device="cuda") * 0.5 + qkv = qkv.to(v.dtype) + q, k, val = qkv[:, 0], qkv[:, 1, : v.num_kv_heads], qkv[:, 2, : v.num_kv_heads] + return q, k, val, np.cumsum([0] + list(lens)) + + +def _assert_non_causal_close(got, q, k, v, what): + """Within twice the error of TRITON_ATTN's mm-encoder kernel against the + float reference, for one sequence.""" + from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd + + m, hq, hd = q.shape + g = hq // k.shape[1] + # A few heads at a time: all 64 heads of a long sequence take 16 GiB. + step = max(1, (1 << 28) // (m * m)) + ref = torch.cat( + [ + torch.nn.functional.scaled_dot_product_attention( + *( + x[:, h : h + step].float().transpose(0, 1) + for x in (q, k.repeat_interleave(g, 1), v.repeat_interleave(g, 1)) + ), + scale=hd**-0.5, + ).transpose(0, 1) + for h in range(0, hq, step) + ], + dim=1, + ) + triton = torch.empty_like(q) + lens = torch.tensor([m], device=q.device, dtype=torch.int32) + context_attention_fwd( + q.contiguous(), k.contiguous(), v.contiguous(), triton, + b_start_loc=torch.zeros_like(lens), b_seq_len=lens, max_input_len=m, + is_causal=False, softmax_scale=hd**-0.5, + ) # fmt: skip + err = (got.float().nan_to_num(1e9) - ref).abs().max().item() + bound = 2 * (triton.float() - ref).abs().max().item() + 1e-4 + assert err <= bound, f"{what}: max error {err:.2e}, Triton's x2 {bound:.2e}" + + +@pytest.mark.parametrize("unit", sorted({_prefill_unit(v) for v in _non_causal_rows()})) +def test_built_in_non_causal_prefill_variants_match_reference(unit): + """Every non-causal build: each token sees every key of its sequence, + read in place, also past the last whole step.""" + _skip_unless_rdna35_or_rdna4() + for v in _non_causal_rows(): + if _prefill_unit(v) != unit: + continue + module = hip.load_prefill(v) + assert module is not None, v + args = (v.num_q_heads, v.num_kv_heads, v.head_size, v.page_size, v.window) + rows = hip.prefill_variants(*args, v.dtype, causal=False) + scratch = hip.make_prefill_scratch(rows, torch.device("cuda")) + lens = [v.min_query_len, v.min_query_len + 37] + if v == rows[0]: + lens += [hip.PREFILL_MIN_M, 100] + for seed, m in enumerate(lens): + q, k, val, _ = _non_causal_inputs(v, [m], seed) + out = torch.full(q.shape, float("nan"), dtype=q.dtype, device=q.device) + seq_lens = torch.tensor([m], device=q.device, dtype=torch.int32) + module.encoder_attn( + q, k, val, out, *scratch, seq_lens, m, v.head_size**-0.5 + ) + torch.accelerator.synchronize() + _assert_non_causal_close(out, q, k, val, f"{v}, L={m}") + + +_NON_CAUSAL_BATCHES = { + "5 x 577": [577] * 5, + "mixed": [1024, 257, 9, 100, 64], + "windows": [64] * 20 + [16, 36, 48], +} + + +@pytest.mark.parametrize("batch", _NON_CAUSAL_BATCHES) +@pytest.mark.parametrize( + "shape", + sorted( + { + (v.num_q_heads, v.num_kv_heads, v.head_size, v.dtype) + for v in _non_causal_rows() + }, + key=str, + ), + ids=str, +) +def test_non_causal_prefill_batch_matches_reference(shape, batch): + """An encoder's segments in one planned launch, each against its + reference.""" + _skip_unless_rdna35_or_rdna4() + hq, hkv, d, dtype = shape + lens = _NON_CAUSAL_BATCHES[batch] + dev = torch.device("cuda") + seq_lens = np.array(lens) + v = hip.prefill_variants(hq, hkv, d, 16, 0, dtype, causal=False)[0] + q, k, val, starts = _non_causal_inputs(v, lens) + v, head_group = hip.prefill_batch_variant_for( + hq, hkv, d, 16, 1, 0, dtype, starts, seq_lens, causal=False + ) + plan, head_group = hip.plan_prefill( + starts, seq_lens, v, 0, hip.prefill_target_workgroups(dev), head_group + ) + scratch = hip.make_prefill_scratch( + hip.prefill_variants(hq, hkv, d, 16, 0, dtype, causal=False), dev + ) + out = torch.full(q.shape, float("nan"), dtype=dtype, device=dev) + hip.load_prefill(v).encoder_attn( + q, k, val, out, *scratch, + torch.tensor(lens, device=dev, dtype=torch.int32), max(lens), d**-0.5, + plan=torch.from_numpy(plan).to(dev), + query_start_loc=torch.from_numpy(starts.astype(np.int32)).to(dev), + head_group=head_group, + ) # fmt: skip + torch.accelerator.synchronize() + for i, n in enumerate(lens): + rows = slice(int(starts[i]), int(starts[i + 1])) + _assert_non_causal_close( + out[rows], q[rows], k[rows], val[rows], f"{batch}: sequence {i} (L={n})" + ) + + +def test_non_causal_rows_are_full_attention(): + """Non-causal builds have no window and run only through encoder_attn; + causal builds only through prefill_attn.""" + rows = _non_causal_rows() + if not rows: + pytest.skip("_rocm_C carries no non-causal RDNA3.5 variants") + assert all(v.window == 0 for v in rows) + _skip_unless_rdna35_or_rdna4() + dev = torch.device("cuda") + causal_sizes = {c.head_size for c in _causal_prefill_rows()} + v = next(v for v in rows if v.head_size in causal_sizes) + scratch = hip.make_prefill_scratch([v], dev) + q, k, val, _ = _non_causal_inputs(v, [64]) + seq_lens = torch.tensor([64], device=dev, dtype=torch.int32) + kv = torch.zeros(4, v.num_kv_heads, 16, 2 * v.head_size, dtype=v.dtype, device=dev) + table = torch.arange(4, dtype=torch.int32, device=dev) + out = torch.empty(q.shape, dtype=q.dtype, device=dev) + with pytest.raises(RuntimeError, match="rdna_encoder_attn"): + hip.load_prefill(v).prefill_attn( + q.contiguous(), kv, table, out, *scratch, seq_lens, 64, 0.1 + ) + causal = next(c for c in _causal_prefill_rows() if c.head_size == v.head_size) + with pytest.raises(RuntimeError, match="rdna_prefill_attn"): + hip.load_prefill(causal).encoder_attn( + q, k, val, out, *scratch, seq_lens, 64, 0.1 + ) + + +@pytest.mark.parametrize( + "segments", + [[4, 8, 1, 3], [577, 5, 1024, 8, 9, 64, 8], [9] * 7, [4096], [5], []], + ids=["short", "mixed", "first step", "one", "one short", "none"], +) +def test_encoder_forward_matches_reference(segments): + """rocm_rdna_encoder: every segment, however short, in one non-causal + launch.""" + _skip_unless_rdna35_or_rdna4() + from vllm.v1.attention.ops import rocm_rdna_encoder + + shapes = sorted( + { + (v.num_q_heads, v.num_kv_heads, v.head_size, v.dtype) + for v in _non_causal_rows() + }, + key=str, + ) + if not shapes: + pytest.skip("_rocm_C carries no non-causal RDNA3.5 variants") + torch.manual_seed(0) + for hq, hkv, d, dtype in shapes: + n = sum(segments) + q, k, v = ( + (torch.randn(n, h, d, device="cuda") * 0.5).to(dtype) + for h in (hq, hkv, hkv) + ) + assert rocm_rdna_encoder.unsupported(q, hkv) is None + out = rocm_rdna_encoder.forward(q, k, v, segments, d**-0.5) + start = 0 + for i, s in enumerate(segments): + rows = slice(start, start + s) + ref = torch.nn.functional.scaled_dot_product_attention( + *(x[rows].float().transpose(0, 1) for x in (q, k, v)), + scale=d**-0.5, + enable_gqa=hq != hkv, + ).transpose(0, 1) + err = (out[rows].float() - ref).abs().max().item() + tol = 2e-2 if dtype == torch.bfloat16 else 4e-3 + assert err <= tol, f"{hq}/{hkv}/{d} {dtype}: segment {i} (L={s}) {err:.2e}" + start += s + + # ---------------------------------------------------------------- prefill batches # (M, S) per sequence. @@ -893,7 +1121,8 @@ def test_prefill_batch_mm_prefix_matches_reference(shape): @pytest.mark.parametrize( - "unit", sorted({_prefill_unit(v) for v in _prefill_rows() if v.page_size == 16}) + "unit", + sorted({_prefill_unit(v) for v in _causal_prefill_rows() if v.page_size == 16}), ) def test_prefill_batch_matches_reference(unit): """A batch in one launch, each sequence against its reference.""" @@ -925,7 +1154,7 @@ def test_prefill_batch_matches_reference(unit): sorted( { _prefill_unit(v) - for v in _prefill_rows() + for v in _causal_prefill_rows() if v.page_size == 16 and v.head_size in hip.SUPPORTED_HEAD_SIZES } ), @@ -934,7 +1163,7 @@ def test_mixed_batch_matches_reference(unit): """Decodes then prefills, including a short one, each against its reference.""" _skip_unless_rdna35_or_rdna4() - v = next(v for v in _prefill_rows() if _prefill_unit(v) == unit) + v = next(v for v in _causal_prefill_rows() if _prefill_unit(v) == unit) batches = [ [(1, 3000 + v.window), (1, 77), (1, 1024), (130, 2000 + v.window), (40, 40)], [(1, 3000 + v.window), (1, 77), (1, 1024), (5, 2000 + v.window)], @@ -1013,25 +1242,26 @@ def test_prefill_batch_configs_are_well_formed(): if not table: pytest.skip("_rocm_C carries no RDNA variants") seen = set() - for (hq, hkv, d, page, window), cells in table.items(): + for (hq, hkv, d, page, window, causal), cells in table.items(): for mean_m, keys, min_query_len, head_group in cells: - key = (hq, hkv, d, page, window, mean_m, keys) + key = (hq, hkv, d, page, window, causal, mean_m, keys) assert key not in seen, key seen.add(key) assert mean_m in hip.PREFILL_BATCH_MEAN_M assert keys in hip.PREFILL_BATCH_KEYS assert head_group == 0 or hkv % head_group == 0 for dtype in DTYPES: - steps = hip.prefill_variants(hq, hkv, d, page, window, dtype) - assert min_query_len in {v.min_query_len for v in steps}, key - for config in {k[:5] for k in seen}: + steps = hip.prefill_variants(hq, hkv, d, page, window, dtype, causal) + if steps or causal: + assert min_query_len in {v.min_query_len for v in steps}, key + for config in {k[:6] for k in seen}: assert (*config, hip.PREFILL_BATCH_MEAN_M[0], 0) in seen, config def test_prefill_batch_cell_lookup(monkeypatch): """A batch takes the cell of the largest steps at or below its features.""" table = { - (32, 8, 128, 16, 0): [(1, 0, 16, 8), (1, 2048, 128, 1), (256, 512, 512, 4)] + (32, 8, 128, 16, 0, 1): [(1, 0, 16, 8), (1, 2048, 128, 1), (256, 512, 512, 4)] } monkeypatch.setattr(hip, "_prefill_batch_table", lambda: table) cell = functools.partial(hip.prefill_batch_cell, 32, 8, 128, 16, 0) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 67b068b68e88..f5d7d37b0841 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -26,6 +26,7 @@ from vllm.triton_utils import HAS_TRITON from vllm.utils import random_uuid from vllm.utils.hashing import safe_hash +from vllm.v1.attention.backends.registry import AttentionBackendEnum from .attention import AttentionConfig, HiSparseConfig from .aux_output import AuxOutputConfig @@ -1458,6 +1459,24 @@ def __post_init__(self): "disable one of them." ) + mm_backend = ( + self.model_config.multimodal_config.mm_encoder_attn_backend + if self.model_config is not None + and self.model_config.multimodal_config is not None + else None + ) + if ( + mm_backend is not None + and mm_backend == AttentionBackendEnum.ROCM_RDNA_ATTN + and self.compilation_config.cudagraph_mm_encoder + ): + raise ValueError( + "--mm-encoder-attn-backend ROCM_RDNA_ATTN is incompatible with " + "cudagraph_mm_encoder=True: its kernels are planned on the host " + "from each call's segment lengths, which a CUDA graph would " + "replay unchanged. Please disable one of them." + ) + self._verify_sampling_replay_config() self._verify_trace_replay_config() diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index dbd71bcc70a5..458ced9a2f6a 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -33,9 +33,11 @@ from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.vit_attn_wrappers import ( + remember_rdna_cu_seqlens, vit_aiter_fp8_attn_wrapper, vit_flash_attn_wrapper, vit_flashinfer_wrapper, + vit_rdna_attn_wrapper, vit_torch_sdpa_wrapper, vit_triton_attn_wrapper, ) @@ -234,6 +236,7 @@ def compute_max_seqlen( AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.ROCM_AITER_FA, AttentionBackendEnum.TRITON_ATTN, + AttentionBackendEnum.ROCM_RDNA_ATTN, AttentionBackendEnum.FLASHINFER, ) and len(cu_seqlens) >= 2 @@ -312,8 +315,10 @@ def maybe_recompute_cu_seqlens( ) cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v]) - cu_seqlens = async_tensor_h2d(cu_seqlens, device=device) - return cu_seqlens + cu_seqlens_gpu = async_tensor_h2d(cu_seqlens, device=device) + if attn_backend == AttentionBackendEnum.ROCM_RDNA_ATTN: + remember_rdna_cu_seqlens(cu_seqlens_gpu, cu_seqlens) + return cu_seqlens_gpu def __init__( self, @@ -675,6 +680,33 @@ def _forward_triton( output = output.reshape(bsz, q_len, -1) return output + def _forward_rdna( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, + ) -> torch.Tensor: + """ROCM_RDNA_ATTN's kernels; TRITON_ATTN's for calls they + cannot serve. Input shapes as _forward_triton.""" + bsz, q_len = query.size()[:2] + kv_len = key.size(1) + is_reshaped = query.dim() != 4 + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len) + output = vit_rdna_attn_wrapper( + q=query, + k=key, + v=value, + batch_size=bsz, + scale=self.scale, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + if is_reshaped: + output = output.reshape(bsz, q_len, -1) + return output + @torch.no_grad() def _record_amax_and_update_scales( self, @@ -843,6 +875,8 @@ def forward_cuda( return self._forward_fa(query, key, value, cu_seqlens, max_seqlen) elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN: return self._forward_triton(query, key, value, cu_seqlens, max_seqlen) + elif self.attn_backend == AttentionBackendEnum.ROCM_RDNA_ATTN: + return self._forward_rdna(query, key, value, cu_seqlens, max_seqlen) elif self.attn_backend == AttentionBackendEnum.FLASHINFER: return self._forward_flashinfer( query, key, value, cu_seqlens, max_seqlen, sequence_lengths diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 2413c20476a2..5562159f67c9 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -698,6 +698,7 @@ def __init__( AttentionBackendEnum.ROCM_AITER_FA, AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.ROCM_RDNA_ATTN, }: raise RuntimeError( f"Qwen2.5-VL does not support {self.attn_backend} backend now." diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 1595386993dd..82ee3b9d587c 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -660,6 +660,7 @@ def __init__( AttentionBackendEnum.ROCM_AITER_FA, AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.ROCM_RDNA_ATTN, }: raise RuntimeError( f"Qwen3-VL does not support {self.attn_backend} backend now." diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index c81bdb3f90ef..7b91bd1521da 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -853,12 +853,18 @@ def get_attn_backend_cls( @classmethod def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]: - return [ + backends = [ AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.ROCM_AITER_FA, AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.TORCH_SDPA, ] + # The RDNA kernels, on request (--mm-encoder-attn-backend); calls they + # cannot serve run on TRITON_ATTN. gfx1152 and gfx1153 share gfx1150's + # tables and gfx1200 gfx1201's, but they are not tested yet. + if on_gfx115x() or on_rdna4(): + backends.append(AttentionBackendEnum.ROCM_RDNA_ATTN) + return backends @classmethod def get_vit_attn_backend( diff --git a/vllm/v1/attention/ops/rocm_rdna_encoder.py b/vllm/v1/attention/ops/rocm_rdna_encoder.py new file mode 100644 index 000000000000..4f6775a79f2c --- /dev/null +++ b/vllm/v1/attention/ops/rocm_rdna_encoder.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCM_RDNA_ATTN for MMEncoderAttention (gfx115x, RDNA4): an +encoder's segments, each token seeing every key of its segment, in one launch +of the non-causal prefill builds, which read Q, K and V in place. +""" + +import functools +from typing import NamedTuple + +import numpy as np +import torch + +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.ops.rocm_rdna_hip import ( + PREFILL_MIN_M, + SUPPORTED_HEAD_SIZES, + PrefillVariant, + load_prefill, + make_prefill_scratch, + plan_prefill, + prefill_batch_variant_for, + prefill_target_workgroups, + prefill_variant_for, + prefill_variants, +) + +PAGE_SIZE = 16 +_HND = 1 +# Per (shape, device): prefill scratch, sized for every non-causal build. +_PREFILL_SCRATCH: dict = {} + + +def unsupported(q: torch.Tensor, num_kv_heads: int) -> str | None: + """Why the kernels cannot serve this call, or None.""" + if q.dtype not in (torch.float16, torch.bfloat16): + return f"dtype {q.dtype}" + hq, d = q.shape[1], q.shape[2] + if d not in SUPPORTED_HEAD_SIZES: + return f"head size {d}" + if hq % num_kv_heads: + return "q heads must divide evenly over kv heads" + if not prefill_variants(hq, num_kv_heads, d, PAGE_SIZE, 0, q.dtype, False): + return f"no non-causal build for {hq}/{num_kv_heads}/{d} {q.dtype}" + return None + + +class _Layout(NamedTuple): + """How one list of segment lengths runs, on the device.""" + + variant: PrefillVariant + seq_lens: torch.Tensor + # The plan and query starts; None for one segment. + plan: torch.Tensor | None + query_starts: torch.Tensor | None + head_group: int + longest: int + + +# A video gives a list per clip shape; the entries are a few small tensors. +@functools.lru_cache(maxsize=128) +def _layout( + seg_lens: tuple[int, ...], + hq: int, + hkv: int, + d: int, + dtype: torch.dtype, + dev: torch.device, +) -> _Layout: + def to_dev(x: np.ndarray) -> torch.Tensor: + return async_tensor_h2d(np.ascontiguousarray(x), dev, torch.int32) + + lens = np.asarray(seg_lens, dtype=np.int64) + shape = (hq, hkv, d, PAGE_SIZE) + if len(lens) == 1: + # The first row also serves fewer than PREFILL_MIN_M tokens. + variant = prefill_variant_for( + *shape, _HND, 0, dtype, max(int(lens[0]), PREFILL_MIN_M), causal=False + ) + assert variant is not None, "unsupported() checks the builds" + return _Layout(variant, to_dev(lens), None, None, 0, int(lens[0])) + starts = np.concatenate([[0], np.cumsum(lens)]) + variant, head_group = prefill_batch_variant_for( + *shape, _HND, 0, dtype, starts, lens, causal=False + ) + assert variant is not None, "unsupported() checks the builds" + plan, head_group = plan_prefill( + starts, lens, variant, 0, prefill_target_workgroups(dev), head_group + ) + return _Layout( + variant, to_dev(lens), to_dev(plan), to_dev(starts), head_group, + int(lens.max()), + ) # fmt: skip + + +def _in_place(x: torch.Tensor) -> torch.Tensor: + """x, or a copy with each token's heads contiguous.""" + return x if x.stride(2) == 1 and x.stride(1) == x.shape[2] else x.contiguous() + + +def forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seg_lens: list[int], + scale: float, +) -> torch.Tensor: + """Attention of each segment's tokens over its keys, no mask. + + Args: + q: (tokens, Hq, D), the segments one after the other. + k: (tokens, Hkv, D). + v: (tokens, Hkv, D). + seg_lens: Each segment's tokens, on the host. + scale: Softmax scale. + + Returns: + (tokens, Hq, D). `unsupported` must have returned None. + + """ + dev, dtype = q.device, q.dtype + out = torch.empty(q.shape, dtype=dtype, device=dev) + if not seg_lens: + return out + hq, d, hkv = q.shape[1], q.shape[2], k.shape[1] + lay = _layout(tuple(seg_lens), hq, hkv, d, dtype, dev) + q, k, v = _in_place(q), _in_place(k), _in_place(v) + prefill = load_prefill(lay.variant) + assert prefill is not None + shape = (hq, hkv, d, PAGE_SIZE, 0, dtype) + scratch = _PREFILL_SCRATCH.get((shape, dev)) + if scratch is None or scratch[1].numel() < 2 * lay.variant.scratch_rows(dev): + scratch = make_prefill_scratch( + [*prefill_variants(*shape, causal=False), lay.variant], dev + ) + _PREFILL_SCRATCH[(shape, dev)] = scratch + if lay.plan is None or lay.plan.shape[0]: + prefill.encoder_attn( + q, k, v, out, *scratch, lay.seq_lens, lay.longest, scale, + plan=lay.plan, query_start_loc=lay.query_starts, + head_group=lay.head_group, + ) # fmt: skip + return out diff --git a/vllm/v1/attention/ops/rocm_rdna_hip.py b/vllm/v1/attention/ops/rocm_rdna_hip.py index 9f2eae8501da..5ba0fa3febfe 100644 --- a/vllm/v1/attention/ops/rocm_rdna_hip.py +++ b/vllm/v1/attention/ops/rocm_rdna_hip.py @@ -268,7 +268,7 @@ def expected_kv_cache_strides(variant: KernelVariant) -> tuple[int, int, int]: MAX_M = 8 # Head dims the decode kernel supports. -SUPPORTED_HEAD_SIZES = (64, 96, 128, 256, 512) +SUPPORTED_HEAD_SIZES = (32, 48, 64, 72, 80, 96, 128, 256, 512) # ---------------------------------------------------------------- prefill @@ -304,6 +304,8 @@ def prefill_target_workgroups(device: torch.device) -> int: "key_waves", "max_segments", "head_group", + "causal", + "row_tiles", ) @@ -326,6 +328,10 @@ class PrefillVariant: max_segments: int head_group: int dtype: torch.dtype = torch.float16 + # 0: every query token sees every key of its sequence; a build of its own. + causal: int = 1 + # 16-row tiles per wave, 1 or (non-causal) 2. + row_tiles: int = 1 @property def name(self) -> str: @@ -335,7 +341,9 @@ def name(self) -> str: @property def block_rows(self) -> int: """Query rows per workgroup.""" - return self.waves // (self.head_dim_split * self.key_waves) * 16 + return ( + self.waves // (self.head_dim_split * self.key_waves) * 16 * self.row_tiles + ) def scratch_rows(self, device: torch.device) -> int: """Most split-KV partial rows a launch may write.""" @@ -361,7 +369,7 @@ def _prefill_builds() -> list[PrefillVariant]: @functools.cache def _prefill_table() -> dict[tuple, list[PrefillVariant]]: - """The builds by M per (Hq, Hkv, D, page, window, dtype).""" + """The builds by M per (Hq, Hkv, D, page, window, dtype, causal).""" table: dict[tuple, list[PrefillVariant]] = {} for v in _prefill_builds(): key = ( @@ -371,6 +379,7 @@ def _prefill_table() -> dict[tuple, list[PrefillVariant]]: v.page_size, v.window, v.dtype, + v.causal, ) table.setdefault(key, []).append(v) for by_m in table.values(): @@ -385,10 +394,12 @@ def prefill_variants( page_size: int, window: int, dtype: torch.dtype, + causal: bool = True, ) -> list[PrefillVariant]: """Every row of one shape, sorted by M.""" return _prefill_table().get( - (num_q_heads, num_kv_heads, head_size, page_size, window, dtype), [] + (num_q_heads, num_kv_heads, head_size, page_size, window, dtype, int(causal)), + [], ) @@ -401,12 +412,13 @@ def prefill_variant_for( window: int, dtype: torch.dtype, num_tokens: int, + causal: bool = True, ) -> PrefillVariant | None: """The row for num_tokens query tokens, or None.""" if layout != _HND or num_tokens < PREFILL_MIN_M: return None rows = prefill_variants( - num_q_heads, num_kv_heads, head_size, page_size, window, dtype + num_q_heads, num_kv_heads, head_size, page_size, window, dtype, causal ) served = [v for v in rows if v.min_query_len <= num_tokens] return served[-1] if served else (rows[0] if rows else None) @@ -428,6 +440,7 @@ def prefill_variant_for( "min_keys", "min_query_len", "head_group", + "causal", ) # The cells' lower bounds. PREFILL_BATCH_MEAN_M = (1, 16, 64, 256, 1024, 4096) @@ -435,7 +448,10 @@ def prefill_variant_for( def prefill_batch_features( - query_start_loc: np.ndarray, seq_lens: np.ndarray, window: int + query_start_loc: np.ndarray, + seq_lens: np.ndarray, + window: int, + causal: bool = True, ) -> tuple[float, float]: """(mean query tokens, mean keys per token) of a batch, padding left out.""" q_lens = np.diff(query_start_loc.astype(np.int64)) @@ -445,7 +461,7 @@ def prefill_batch_features( tokens = int(q_lens.sum()) if not tokens: return 0.0, 0.0 - keys = seq_lens - q_lens + (q_lens + 1) / 2 + keys = seq_lens - q_lens + (q_lens + 1) / 2 if causal else seq_lens if window: keys = np.minimum(keys, window) return tokens / q_lens.shape[0], float((q_lens * keys).sum() / tokens) @@ -453,7 +469,7 @@ def prefill_batch_features( @functools.cache def _prefill_batch_table() -> dict[tuple, list[tuple[int, int, int, int]]]: - """Per (Hq, Hkv, D, page, window), its cells (min mean M, min keys, + """Per (Hq, Hkv, D, page, window, causal), its cells (min mean M, min keys, min_query_len, head group).""" table: dict[tuple, list[tuple[int, int, int, int]]] = {} for f in _built_fields("rdna_prefill_batch_configs", PREFILL_BATCH_CONFIG_FIELDS): @@ -463,6 +479,7 @@ def _prefill_batch_table() -> dict[tuple, list[tuple[int, int, int, int]]]: f["head_size"], f["page_size"], f["window"], + f["causal"], ) table.setdefault(key, []).append( (f["min_mean_m"], f["min_keys"], f["min_query_len"], f["head_group"]) @@ -478,10 +495,11 @@ def prefill_batch_cell( window: int, mean_m: float, keys: float, + causal: bool = True, ) -> tuple[int, int] | None: """(min_query_len, head group) of the cell a batch falls in, or None.""" cells = _prefill_batch_table().get( - (num_q_heads, num_kv_heads, head_size, page_size, window) + (num_q_heads, num_kv_heads, head_size, page_size, window, int(causal)) ) if not cells: return None @@ -503,15 +521,16 @@ def prefill_batch_variant_for( dtype: torch.dtype, query_start_loc: np.ndarray, seq_lens: np.ndarray, + causal: bool = True, ) -> tuple[PrefillVariant | None, int]: """The build and head group for a prefill batch: its cell's, else the build of its tokens.""" - mean_m, keys = prefill_batch_features(query_start_loc, seq_lens, window) + mean_m, keys = prefill_batch_features(query_start_loc, seq_lens, window, causal) cell = prefill_batch_cell( - num_q_heads, num_kv_heads, head_size, page_size, window, mean_m, keys + num_q_heads, num_kv_heads, head_size, page_size, window, mean_m, keys, causal ) rows = prefill_variants( - num_q_heads, num_kv_heads, head_size, page_size, window, dtype + num_q_heads, num_kv_heads, head_size, page_size, window, dtype, causal ) if cell is not None and layout == _HND: row = next((v for v in rows if v.min_query_len == cell[0]), None) @@ -527,6 +546,7 @@ def prefill_batch_variant_for( window, dtype, max(tokens, PREFILL_MIN_M), + causal, ) return v, 0 @@ -554,7 +574,7 @@ def plan_prefill( query_start_loc: The batch's query starts. seq_lens: Each sequence's length, or an upper bound. variant: The build that runs the launch. - window: Sliding window in keys, 0 for full causal. + window: Sliding window in keys, 0 for full causal (or non-causal). target_workgroups: One per SIMD (prefill_target_workgroups). head_group: Kv heads to interleave (0: the build's). @@ -576,7 +596,8 @@ def plan_prefill( tile = np.repeat(starts + tiles - 1, tiles) - np.arange(seq.shape[0]) first = tile * rows // g last = np.minimum(q_lens[seq] - 1, ((tile + 1) * rows - 1) // g) - key_end = context[seq] + last + 1 + # Non-causal: every token sees its whole sequence. + key_end = context[seq] + last + 1 if variant.causal else context[seq] + q_lens[seq] key_begin = ( np.maximum(0, context[seq] + first - window + 1) // key_tile * key_tile if window @@ -654,6 +675,29 @@ def prefill_attn( head_group=head_group, ) + def encoder_attn( + self, + q, + k, + v, + out, + *scratch_and_args, + plan=None, + query_start_loc=None, + head_group=0, + ): + torch.ops._rocm_C.rdna_encoder_attn( + self.index, + q, + k, + v, + out, + *scratch_and_args, + plan=plan, + query_start_loc=query_start_loc, + head_group=head_group, + ) + _prefill_loaded: dict[PrefillVariant, _BuiltinPrefill | None] = {} diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index 676026ad97d2..5ed70021bcb6 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -11,17 +11,26 @@ To use these ops, you must have a recent version of PyTorch installed (>= 2.4.0) """ +import weakref from typing import Any import einops +import numpy as np import torch import torch.nn.functional as F from vllm._aiter_ops import rocm_aiter_ops +from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.utils import create_attention_profiler_scope +from vllm.v1.attention.ops import rocm_rdna_encoder +from vllm.v1.utils import ( + create_attention_profiler_scope, + profiler_scopes_enabled, +) + +logger = init_logger(__name__) def flash_attn_maxseqlen_wrapper( @@ -210,9 +219,15 @@ def triton_attn_wrapper( head_size = q.shape[2] num_kv_heads = k.shape[1] - # Compute segment lengths for block-diagonal attention in multi-segment batches. + # Segment lengths for the profiler label of multi-segment batches: a host + # read, so only when a scope records it and not under CUDA graph capture, + # which forbids it. block_diagonal_segments: list[int] | None = None - if cu_seqlens.numel() > 2: + if ( + cu_seqlens.numel() > 2 + and profiler_scopes_enabled() + and not torch.cuda.is_current_stream_capturing() + ): seg_lens = cu_seqlens[1:] - cu_seqlens[:-1] block_diagonal_segments = seg_lens.tolist() @@ -286,6 +301,100 @@ def vit_triton_attn_wrapper( ) +# Segment lengths per live cu_seqlens tensor, which nothing changes in place: +# the host lengths MMEncoderAttention.maybe_recompute_cu_seqlens made it from, +# else read back from the device once. An entry goes with its tensor. +_RDNA_SEG_LENS: dict[int, list[int]] = {} + + +def _remember_rdna_seg_lens(cu_seqlens: torch.Tensor, seg_lens: list[int]) -> None: + key = id(cu_seqlens) + if key not in _RDNA_SEG_LENS: + weakref.finalize(cu_seqlens, _RDNA_SEG_LENS.pop, key, None) + _RDNA_SEG_LENS[key] = seg_lens + + +def remember_rdna_cu_seqlens(cu_seqlens: torch.Tensor, host: np.ndarray) -> None: + """Note the host lengths of cu_seqlens.""" + _remember_rdna_seg_lens(cu_seqlens, np.diff(host).tolist()) + + +def _rdna_segment_lengths(cu_seqlens: torch.Tensor) -> list[int]: + seg_lens = _RDNA_SEG_LENS.get(id(cu_seqlens)) + if seg_lens is None: + # The plan follows the host, as Triton's segment lengths do. + with gpu_sync_allowed(): + bounds = cu_seqlens.tolist() + seg_lens = [b - a for a, b in zip(bounds, bounds[1:])] + _remember_rdna_seg_lens(cu_seqlens, seg_lens) + return seg_lens + + +def rdna_attn_wrapper( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + batch_size: int, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, +) -> torch.Tensor: + """ROCM_RDNA_ATTN's kernels, else triton_attn_wrapper.""" + # VllmConfig refuses cudagraph_mm_encoder with this backend: its launches + # follow the host's segment lengths, which a graph would replay unchanged. + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("ROCM_RDNA_ATTN cannot run under CUDA graph capture") + + if scale is None: + scale = q.shape[-1] ** -0.5 + qf, kf, vf = (einops.rearrange(x, "b s ... -> (b s) ...") for x in (q, k, v)) + why = rocm_rdna_encoder.unsupported(qf, kf.shape[1]) + if why is not None: + logger.info_once("ROCM_RDNA_ATTN mm-encoder call on TRITON_ATTN: %s", why) + return triton_attn_wrapper(q, k, v, batch_size, scale, cu_seqlens, max_seqlen) + seg_lens = ( + [q.size(1)] * batch_size + if cu_seqlens is None + else _rdna_segment_lengths(cu_seqlens) + ) + with create_attention_profiler_scope( + backend_name="ROCM_RDNA_ATTN", + batch_size=batch_size, + max_query_len=max(seg_lens, default=0), + max_seq_len=max(seg_lens, default=0), + num_heads=qf.shape[1], + head_size=qf.shape[2], + num_kv_heads=kf.shape[1], + dtype=q.dtype, + is_causal=False, + v_token_stride=vf.stride(0), + block_diagonal_segments=seg_lens if len(seg_lens) > 1 else None, + ): + output = rocm_rdna_encoder.forward(qf, kf, vf, seg_lens, scale) + return einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size) + + +direct_register_custom_op( + op_name="rdna_attn_wrapper", + op_func=rdna_attn_wrapper, + fake_impl=triton_attn_wrapper_fake, +) + + +def vit_rdna_attn_wrapper( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + batch_size: int, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, +) -> torch.Tensor: + return torch.ops.vllm.rdna_attn_wrapper( + q, k, v, batch_size, scale, cu_seqlens, max_seqlen + ) + + def apply_sdpa( q: torch.Tensor, k: torch.Tensor, diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index c92cf6130ef6..68acc1278249 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -779,6 +779,14 @@ def report_usage_stats( _PROFILER_FUNC = None +def profiler_scopes_enabled() -> bool: + """Whether record_function_or_nullcontext opens a scope, so a label that + costs work to build is worth building.""" + return bool( + envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING or envs.VLLM_NVTX_SCOPES_FOR_PROFILING + ) + + def record_function_or_nullcontext(name: str) -> AbstractContextManager: global _PROFILER_FUNC