diff --git a/.github/workflows/amd_tests.yml b/.github/workflows/amd_tests.yml index 50e93d965e0..53e52b6c5de 100644 --- a/.github/workflows/amd_tests.yml +++ b/.github/workflows/amd_tests.yml @@ -57,7 +57,12 @@ jobs: - name: Build run: | python setup.py install - - name: Test + - name: AMD Kernel Tests run: | + pytest flash_attn/flash_attn_triton_kernel_decode_amd.py::test_op_fwd + pytest flash_attn/flash_attn_triton_kernel_prefill_amd.py + - name: Flash Attention Tests + run: | + pytest tests/test_flash_attn.py::test_flash_attn_kvcache pytest tests/test_flash_attn.py::test_flash_attn_output - pytest tests/test_flash_attn.py::test_flash_attn_varlen_output \ No newline at end of file + pytest tests/test_flash_attn.py::test_flash_attn_varlen_output diff --git a/.gitignore b/.gitignore index f6af83f262b..2ee76860d53 100644 --- a/.gitignore +++ b/.gitignore @@ -26,9 +26,10 @@ var/ # Dev venv -# Other +# AMD .eggs .vscode core scripts log* +*csv \ No newline at end of file diff --git a/benchmarks/benchmark_flash_attention.py b/benchmarks/benchmark_flash_attention.py index 341ae4b2139..2a2ff2bda24 100644 --- a/benchmarks/benchmark_flash_attention.py +++ b/benchmarks/benchmark_flash_attention.py @@ -2,16 +2,20 @@ # pip install "git+https://github.com/openai/triton.git#egg=triton&subdirectory=python" import pickle import math +import argparse +import sys +import os +import time +import random import torch import torch.nn as nn import torch.nn.functional as F - from einops import rearrange, repeat - from flash_attn.utils.benchmark import benchmark_all, benchmark_forward, benchmark_backward from flash_attn.utils.benchmark import benchmark_fwd_bwd, benchmark_combined from flash_attn import flash_attn_qkvpacked_func +import multiprocessing as mp try: from triton.ops.flash_attention import attention as attention_triton @@ -23,6 +27,10 @@ except ImportError: xops = None +try: + import pandas as pd +except ImportError: + pd = None def flops(batch, seqlen, headdim, nheads, causal, mode="fwd"): assert mode in ["fwd", "bwd", "fwd_bwd"] @@ -32,7 +40,6 @@ def flops(batch, seqlen, headdim, nheads, causal, mode="fwd"): def efficiency(flop, time): return (flop / time / 10**12) if not math.isnan(time) else 0.0 - def attention_pytorch(qkv, dropout_p=0.0, causal=True): """ Arguments: @@ -66,115 +73,201 @@ def time_fwd_bwd(func, *args, **kwargs): time_f, time_b = benchmark_fwd_bwd(func, *args, **kwargs) return time_f[1].mean, time_b[1].mean +def time_fwd(func, *args, **kwargs): + time_f = benchmark_forward(func, *args, **kwargs) + return time_f[1].mean + +def time_bwd(func, *args, **kwargs): + time_b = benchmark_backward(func, *args, **kwargs) + return time_b[1].mean +def run_single_benchmark(args_dict, config, methods, return_dict): + # Convert args_dict back to an object + class Args: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + args = Args(**args_dict) + + os.environ['HIP_VISIBLE_DEVICES'] = str(args.gpu) + torch.cuda.empty_cache() + + causal, headdim, batch_size, seqlen, nheads = config + + print(f"Warming up: {config}") + device = torch.device(args.device) + qkv = torch.randn(batch_size, seqlen, 3, nheads, headdim, device=device, dtype=getattr(torch, args.dtype), + requires_grad=True) + for _ in range(5): # 5 warm-up iterations + for method in methods: + if method == "Flash2": + flash_attn_qkvpacked_func(qkv, args.dropout_p, causal=causal) + elif method == "Pytorch": + attention_pytorch(qkv, args.dropout_p, causal=causal) + elif method == "Triton" and attention_triton is not None: + q, k, v = [x.reshape(batch_size, nheads, seqlen, headdim) for x in qkv.unbind(dim=2)] + attention_triton(q, k, v, causal, headdim**(-0.5), False) + elif method in ["xformers.c", "xformers.f"] and xops is not None: + q, k, v = [x.reshape(batch_size, seqlen, nheads, headdim) for x in qkv.unbind(dim=2)] + op = (xops.fmha.cutlass.FwOp, xops.fmha.cutlass.BwOp) if method == "xformers.c" else (xops.fmha.flash.FwOp, xops.fmha.flash.BwOp) + xops.memory_efficient_attention(q, k, v, attn_bias=xops.LowerTriangularMask() if causal else None, op=op) + + results = {} + for method in methods: + if method == "Flash2": + if args.mode in ["fwd_bwd", "fwd"]: + time_f = time_fwd(flash_attn_qkvpacked_func, qkv, args.dropout_p, causal=causal, repeats=args.repeats, verbose=False) + results[f'{method} fwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd"), time_f) + if args.mode in ["fwd_bwd", "bwd"]: + time_b = time_bwd(flash_attn_qkvpacked_func, qkv, args.dropout_p, causal=causal, repeats=args.repeats, verbose=False) + results[f'{method} bwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="bwd"), time_b) + elif method == "Pytorch": + if args.mode in ["fwd_bwd", "fwd"]: + time_f = time_fwd(attention_pytorch, qkv, args.dropout_p, causal=causal, repeats=args.repeats, verbose=False) + results[f'{method} fwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd"), time_f) + if args.mode in ["fwd_bwd", "bwd"]: + time_b = time_bwd(attention_pytorch, qkv, args.dropout_p, causal=causal, repeats=args.repeats, verbose=False) + results[f'{method} bwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="bwd"), time_b) + elif method == "Triton" and attention_triton is not None: + q, k, v = [x.reshape(batch_size, nheads, seqlen, headdim) for x in qkv.unbind(dim=2)] + if args.mode in ["fwd_bwd", "fwd"]: + time_f = time_fwd(attention_triton, q, k, v, causal, headdim**(-0.5), False, repeats=args.repeats, verbose=False) + results[f'{method} fwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd"), time_f) + if args.mode in ["fwd_bwd", "bwd"]: + time_b = time_bwd(attention_triton, q, k, v, causal, headdim**(-0.5), False, repeats=args.repeats, verbose=False) + results[f'{method} bwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="bwd"), time_b) + elif method in ["xformers.c", "xformers.f"] and xops is not None: + q, k, v = [x.reshape(batch_size, seqlen, nheads, headdim) for x in qkv.unbind(dim=2)] + op = (xops.fmha.cutlass.FwOp, xops.fmha.cutlass.BwOp) if method == "xformers.c" else (xops.fmha.flash.FwOp, xops.fmha.flash.BwOp) + if args.mode in ["fwd_bwd", "fwd"]: + time_f = time_fwd(xops.memory_efficient_attention, q, k, v, attn_bias=xops.LowerTriangularMask() if causal else None, op=op) + results[f'{method} fwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd"), time_f) + if args.mode in ["fwd_bwd", "bwd"]: + time_b = time_bwd(xops.memory_efficient_attention, q, k, v, attn_bias=xops.LowerTriangularMask() if causal else None, op=op) + results[f'{method} bwd (TFLOPs/s)'] = efficiency(flops(batch_size, seqlen, headdim, nheads, causal, mode="bwd"), time_b) + + return_dict[config] = results + +def run_benchmark(args): + print("args:", args) + device = args.device + dtype = getattr(torch, args.dtype) + + methods = args.methods + if attention_triton is None and "Triton" in methods: + methods.remove("Triton") + if xops is None and ("xformers.c" in methods or "xformers.f" in methods): + methods = [m for m in methods if not m.startswith("xformers")] + + configs = [(causal, headdim, batch_size, seqlen, nheads) + for causal in args.causal + for headdim, nheads in zip(args.headdim, args.nheads) + for batch_size, seqlen in zip(args.batch_size, args.seqlen)] + + all_results = [] + for pass_num in range(args.passes): + print(f"Starting pass {pass_num + 1}/{args.passes}") + random.shuffle(configs) + for config in configs: + manager = mp.Manager() + return_dict = manager.dict() + args_dict = vars(args) # Convert args to a dictionary + p = mp.Process(target=run_single_benchmark, args=(args_dict, config, methods, return_dict)) + p.start() + p.join() + + if config not in return_dict: + print(f"Error: No results returned for config {config}") + continue + + results = return_dict[config] + causal, headdim, batch_size, seqlen, nheads = config + row = {'Pass': pass_num + 1, 'Batch Size': batch_size, 'Seq Len': seqlen, "Num of Heads": nheads, 'Dim of Head': headdim, 'Causal': causal} + row.update(results) + all_results.append(row) + + print(f"Completed: {config}") + print(results) + + time.sleep(args.cooldown) + + df = pd.DataFrame(all_results) + print(df.to_string(index=False)) + csv_filename = f'benchmark_results_{args.mode}.csv' + df.to_csv(csv_filename, index=False) + print(f"\nResults saved to {csv_filename}") + + print("\nAveraged Results:") + df_avg = df.drop(columns=["Pass"]).groupby(['Batch Size', 'Seq Len', "Num of Heads", 'Dim of Head', 'Causal']).mean().reset_index() + print(df_avg.to_string(index=False)) + csv_filename = f'benchmark_avg_results_{args.mode}.csv' + df_avg.to_csv(csv_filename, index=False) + print(f"\nAvg Results saved to {csv_filename}") + + + return df + +if __name__ == "__main__": + mp.set_start_method('spawn') # Set the start method to 'spawn' + + parser = argparse.ArgumentParser(description="Run attention benchmark with custom configurations.") + parser.add_argument("--methods", nargs="+", default=["Flash2", "Pytorch"], + help="Attention methods to benchmark") + parser.add_argument("--mode", choices=["fwd", "bwd", "fwd_bwd"], default="fwd", + help="Benchmark mode: forward, backward, or both") + parser.add_argument("--device", default="cuda", help="Device to run the benchmark on") + parser.add_argument("--dtype", default="float16", choices=["float16", "float32", "bfloat16"], + help="Data type for the tensors") + parser.add_argument("--causal", nargs="+", type=lambda x: x.lower() == 'true', default=[False, True], + help="Whether to use causal attention") + parser.add_argument("--dim", type=int, default=2048, help="Total dimension of the model") + parser.add_argument("--nheads", nargs="+", type=int, help="Number of attention heads") + parser.add_argument("--headdim", nargs="+", type=int, default=[64, 128], help="Dimension(s) of each attention head") + parser.add_argument("--batch_size", nargs="+", type=int, default=[32, 16, 8, 4, 2, 1], + help="Batch sizes to benchmark") + parser.add_argument("--seqlen", nargs="+", type=int, default=[512, 1024, 2048, 4096, 8192, 16384], + help="Sequence lengths to benchmark") + parser.add_argument("--dropout_p", type=float, default=0.0, help="Dropout probability") + parser.add_argument("--repeats", type=int, default=30, help="Number of repetitions for each benchmark") + parser.add_argument("--passes", type=int, default=3, help="Number of passes through all configurations") + parser.add_argument("--cooldown", type=float, default=2.0, help="Cool-down time between configurations (seconds)") + parser.add_argument("--gpu", type=int, default=0, help="GPU device to use") + + args = parser.parse_args() + + # Check if all three parameters are provided + if args.dim != parser.get_default('dim') and args.nheads is not None and args.headdim != parser.get_default('headdim'): + print("Error: You have provided values for dim, nheads, and headdim. Please provide at most two of these parameters.") + sys.exit(1) + + # Validate and adjust arguments + if args.nheads is not None: + if args.dim != parser.get_default('dim'): + print(f"Running benchmark with dim={args.dim}, nheads={args.nheads}") + args.headdim = [args.dim // nh for nh in args.nheads] + print(f"Calculated headdim: {args.headdim}") + # Check if any calculated headdim is 0 + if any(hd == 0 for hd in args.headdim): + print("Error: Some number of heads are larger than the total dimension. Please adjust your input.") + sys.exit(1) + else: + print(f"Running benchmark with nheads={args.nheads}, headdim={args.headdim}") + args.dim = max(nh * hd for nh, hd in zip(args.nheads, args.headdim)) + elif args.dim != parser.get_default('dim'): + print(f"Running benchmark with dim={args.dim}, headdim={args.headdim}") + args.nheads = [args.dim // hd for hd in args.headdim] + print(f"Calculated nheads: {args.nheads}") + # Check if any calculated nheads is 0 + if any(nh == 0 for nh in args.nheads): + print("Error: Some head dimensions are larger than the total dimension. Please adjust your input.") + sys.exit(1) + else: + print(f"Running benchmark with default dim={args.dim}, headdim={args.headdim}") + args.nheads = [args.dim // hd for hd in args.headdim] + print(f"Calculated nheads: {args.nheads}") + + # Ensure nheads and headdim have the same length + if len(args.nheads) != len(args.headdim): + print("Error: The number of values for nheads and headdim must be the same.") + sys.exit(1) -repeats = 30 -device = 'cuda' -dtype = torch.float16 - -bs_seqlen_vals = [(32, 512), (16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)] -causal_vals = [False, True] -headdim_vals = [64, 128] -dim = 2048 -dropout_p = 0.0 - -methods = (["Flash2", "Pytorch"] - + (["Triton"] if attention_triton is not None else []) - + (["xformers.c"] if xops is not None else []) - + (["xformers.f"] if xops is not None else [])) - -time_f = {} -time_b = {} -time_f_b = {} -speed_f = {} -speed_b = {} -speed_f_b = {} -for causal in causal_vals: - for headdim in headdim_vals: - for batch_size, seqlen in bs_seqlen_vals: - config = (causal, headdim, batch_size, seqlen) - nheads = dim // headdim - qkv = torch.randn(batch_size, seqlen, 3, nheads, headdim, device=device, dtype=dtype, - requires_grad=True) - f, b = time_fwd_bwd( - flash_attn_qkvpacked_func, qkv, dropout_p, causal=causal, repeats=repeats, verbose=False - ) - time_f[config, "Flash2"] = f - time_b[config, "Flash2"] = b - - try: - qkv = qkv.detach().requires_grad_(True) - f, b = time_fwd_bwd( - attention_pytorch, qkv, dropout_p, causal=causal, repeats=repeats, verbose=False - ) - except: # Skip if OOM - f, b = float('nan'), float('nan') - time_f[config, "Pytorch"] = f - time_b[config, "Pytorch"] = b - - if attention_triton is not None: - q, k, v = [torch.randn(batch_size, nheads, seqlen, headdim, device=device, dtype=dtype, - requires_grad=True) for _ in range(3)] - # Try both values of sequence_parallel and pick the faster one - try: - f, b = time_fwd_bwd( - attention_triton, q, k, v, causal, headdim**(-0.5), - False, repeats=repeats, verbose=False - ) - except: - f, b = float('nan'), float('inf') - try: - _, b0 = time_fwd_bwd( - attention_triton, q, k, v, causal, headdim**(-0.5), - True, repeats=repeats, verbose=False - ) - except: - b0 = float('inf') - time_f[config, "Triton"] = f - time_b[config, "Triton"] = min(b, b0) if min(b, b0) < float('inf') else float('nan') - - if xops is not None: - q, k, v = [torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype, - requires_grad=True) for _ in range(3)] - f, b = time_fwd_bwd( - xops.memory_efficient_attention, q, k, v, - attn_bias=xops.LowerTriangularMask() if causal else None, - op=(xops.fmha.cutlass.FwOp, xops.fmha.cutlass.BwOp) - ) - time_f[config, "xformers.c"] = f - time_b[config, "xformers.c"] = b - - if xops is not None: - q, k, v = [torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype, - requires_grad=True) for _ in range(3)] - f, b = time_fwd_bwd( - xops.memory_efficient_attention, q, k, v, - attn_bias=xops.LowerTriangularMask() if causal else None, - op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp) - ) - time_f[config, "xformers.f"] = f - time_b[config, "xformers.f"] = b - - print(f"### causal={causal}, headdim={headdim}, batch_size={batch_size}, seqlen={seqlen} ###") - for method in methods: - time_f_b[config, method] = time_f[config, method] + time_b[config, method] - speed_f[config, method] = efficiency( - flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd"), - time_f[config, method] - ) - speed_b[config, method] = efficiency( - flops(batch_size, seqlen, headdim, nheads, causal, mode="bwd"), - time_b[config, method] - ) - speed_f_b[config, method] = efficiency( - flops(batch_size, seqlen, headdim, nheads, causal, mode="fwd_bwd"), - time_f_b[config, method] - ) - print( - f"{method} fwd: {speed_f[config, method]:.2f} TFLOPs/s, " - f"bwd: {speed_b[config, method]:.2f} TFLOPs/s, " - f"fwd + bwd: {speed_f_b[config, method]:.2f} TFLOPs/s" - ) - - -# with open('flash2_attn_time.plk', 'wb') as fp: -# pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL) + results = run_benchmark(args) \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_interface_amd.py b/flash_attn/flash_attn_triton_interface_amd.py index 2ca74996807..f74cc7860af 100644 --- a/flash_attn/flash_attn_triton_interface_amd.py +++ b/flash_attn/flash_attn_triton_interface_amd.py @@ -1,8 +1,9 @@ -from .flash_attn_triton_kernel_amd import MetaData, attention, get_shape_from_layout, _attn_bwd_preprocess, _attn_bwd import torch import triton +from .flash_attn_triton_kernel_prefill_amd import MetaData, attention_prefill, get_shape_from_layout, _attn_bwd_preprocess, _attn_bwd +from .flash_attn_triton_kernel_decode_amd import attention_decode -DEBUG=False +DEBUG = False def fwd(q, k, @@ -31,7 +32,7 @@ def fwd(q, print("gen_:", gen_) if dropout_p != 0.0: - raise ValueError("dropout is not supported on HIP") + raise ValueError("dropout is not supported on AMD yet") if o is None: o = torch.empty_like(q) @@ -60,7 +61,7 @@ def fwd(q, input_metadata.check_args(q, k, v, o) # Perform the forward attention computation - tri_out, encoded_softmax = attention(q, k, v, o, input_metadata) + tri_out, encoded_softmax = attention_prefill(q, k, v, o, input_metadata) softmax_lse = encoded_softmax softmax_p = encoded_softmax @@ -93,9 +94,23 @@ def varlen_fwd( print("q:", q.shape) print("k:", k.shape) print("v:", v.shape) + print("cu_seqlens_q:", cu_seqlens_q) + print("cu_seqlens_k:", cu_seqlens_k) + print("block_table_:", block_table_) + print("alibi_slopes:", alibi_slopes) + print("max_seqlen_q:", max_seqlen_q) + print("max_seqlen_k:", max_seqlen_k) + print("dropout_p:", dropout_p) + print("softmax_scale:", softmax_scale) + print("zero_tensors:", zero_tensors) + print("causal:", causal) + print("window_size_left:", window_size_left) + print("window_size_right:", window_size_right) + print("return_softmax:", return_softmax) + print("gen_:", gen_) if dropout_p != 0.0: - raise ValueError("dropout is not supported on HIP") + raise ValueError("dropout is not supported on AMD yet") if o is None: o = torch.empty_like(q) @@ -123,14 +138,14 @@ def varlen_fwd( input_metadata.check_args(q, k, v, o) # Perform the forward attention computation - tri_out, encoded_softmax = attention(q, k, v, o, input_metadata) + tri_out, encoded_softmax = attention_prefill(q, k, v, o, input_metadata) softmax_lse = encoded_softmax softmax_p = encoded_softmax return tri_out, q , k , v, o, softmax_lse, softmax_p, torch.get_rng_state() -def fwd_kvcache( +def fwd_kvcache( q, k_cache, v_cache, @@ -149,150 +164,67 @@ def fwd_kvcache( window_size_right, rotary_interleaved, num_splits): - pass - - -def bwd(dout, q, k, v, out, softmax_lse, dq, dk, dv, alibi_slopes, dropout_p, softmax_scale, causal, window_size_left, - window_size_right, deterministic, gen_, rng_state): + if DEBUG: - print("flash_attn_triton_amd.py::bwd") - print("q:", q.shape) - print("k:", k.shape) - print("v:", v.shape) - print("softmax_lse:", softmax_lse) - print("dq:", dq.shape) - print("dk:", dk.shape) - print("dv:", dv.shape) + print() + print("flash_attn_triton_amd.py::fwd_kvcache") + print("q:", q, q.shape) + print("k_cache:", k_cache, k_cache.shape) + print("v_cache:", v_cache, v_cache.shape) + print("k:", k, k.shape if k is not None else None) + print("v:", v, v.shape if v is not None else None) + print("cache_seqlens:", cache_seqlens, cache_seqlens.size()) + print("rotary_cos:", rotary_cos) + print("rotary_sin:", rotary_sin) + print("cache_batch_idx:", cache_batch_idx) + print("block_table:", block_table, block_table.shape if block_table is not None else None) print("alibi_slopes:", alibi_slopes) - print("dropout_p:", dropout_p) + print("out:", out) print("softmax_scale:", softmax_scale) print("causal:", causal) print("window_size_left:", window_size_left) print("window_size_right:", window_size_right) - print("deterministic:", deterministic) - print("gen_:", gen_) - print("rng_state:", rng_state) - + print("rotary_interleaved:", rotary_interleaved) + print("num_splits:", num_splits) + if out is None: out = torch.empty_like(q) - # Ensure the tensors have requires_grad=True - q.requires_grad_() - k.requires_grad_() - v.requires_grad_() - out.requires_grad_() - - # Create metadata object - metadata = MetaData(sm_scale=softmax_scale) - metadata.max_seqlens_q = q.shape[1] - metadata.max_seqlens_k = k.shape[1] - metadata.layout = "bshd" + # fill metadata + input_metadata = MetaData(sm_scale=softmax_scale) + input_metadata.layout = "bshd" + input_metadata.max_seqlens_q = q.shape[1] + input_metadata.max_seqlens_k = k_cache.shape[1] + input_metadata.cache_seqlens = cache_seqlens + input_metadata.cache_batch_idx = cache_batch_idx - if metadata == 'bshd': - q = q.transpose(1, 2).clone() - k = k.transpose(1, 2).clone() - v = v.transpose(1, 2).clone() + if k is not None and v is not None: + input_metadata.new_kv = True + input_metadata.seqlen_new = k.shape[1] + input_metadata.k_new = k + input_metadata.v_new = v - batch = q.shape[0] - nheads_q = q.shape[1] - BLOCK_DMODEL = q.shape[3] - - # Setup metadata if causal: - metadata.need_causal() + input_metadata.need_causal() - # if bias is not None: - # metadata.need_bias(bias, q.shape[0], q.shape[1], q.shape[2], k.shape[2]) - - return_softmax = True if alibi_slopes is not None: - metadata.need_alibi(alibi_slopes, batch, nheads_q) - - if dropout_p > 0.0: - metadata.need_dropout(dropout_p, return_softmax) + batch, _ , nheads_q, _= q.shape + input_metadata.need_alibi(alibi_slopes, batch, nheads_q) - # Check arguments - metadata.check_args(q, k, v, out) - - # write your own version backward - M = torch.empty((batch, nheads_q, metadata.max_seqlens_q), device=q.device, dtype=torch.float32) # this passed from + # launch kernel + tri_out = attention_decode(q, k_cache, v_cache, input_metadata) - if torch.version.hip is not None: - BLOCK = 64 - else: - BLOCK = 128 - o = out - do = dout - sm_scale = softmax_scale - assert do.is_contiguous() - assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() - seqlen_q = q.shape[2] - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - BATCH, N_CTX, N_HEAD = q.shape[:3] - PRE_BLOCK = 128 - # NUM_WARPS, NUM_STAGES = 4, 1 - BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32 - BLK_SLICE_FACTOR = 2 - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) - arg_k = k - arg_k = arg_k * (sm_scale * RCP_LN2) if DEBUG: - print("N_CTX:", N_CTX) - # assert N_CTX % PRE_BLOCK == 0 + print() + print("tri_out:", tri_out, tri_out.shape) + + return tri_out, None - delta = torch.empty_like(M) - _, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1] - # padded_head = (Lk != ctx.BLOCK_DMODEL) - grid_preprocess = (triton.cdiv(do.shape[2], BLOCK), do.shape[1], do.shape[0]) - _attn_bwd_preprocess[grid_preprocess]( - o, - do, - delta, - o.stride(0), - o.stride(1), - o.stride(2), - o.stride(3), - do.stride(0), - do.stride(1), - do.stride(2), - do.stride(3), - seqlen_q, - head_dim=Lk, - BLOCK_M=BLOCK, - D_HEAD=BLOCK_DMODEL, - ) - grid = lambda META: (triton.cdiv(N_CTX, META['BLOCK_N1']), 1, BATCH * N_HEAD) - _attn_bwd[grid]( - q, - arg_k, - v, - sm_scale, - alibi_slopes, - do, - dq, - dk, - dv, - M, - delta, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - N_HEAD, - N_CTX, - BLOCK_DMODEL= BLOCK_DMODEL, - BLOCK_M1=BLOCK_M1, - BLOCK_N1=BLOCK_N1, - BLOCK_M2=BLOCK_M2, - BLOCK_N2=BLOCK_N2, - BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, - USE_ALIBI=False if alibi_slopes is None else True, - ) - return dq, dk, dv, None +def bwd(dout, q, k, v, out, softmax_lse, dq, dk, dv, alibi_slopes, dropout_p, softmax_scale, causal, window_size_left, + window_size_right, deterministic, gen_, rng_state): + raise ValueError("bwd is not supported on AMD yet") def varlen_bwd(dout, q, k, v, out, softmax_lse, dq, dk, dv, *args, **kwargs): - pass \ No newline at end of file + raise ValueError("varlen_bwd is not supported on AMD yet") \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_kernel_decode_amd.py b/flash_attn/flash_attn_triton_kernel_decode_amd.py new file mode 100644 index 00000000000..7bf4586641e --- /dev/null +++ b/flash_attn/flash_attn_triton_kernel_decode_amd.py @@ -0,0 +1,1058 @@ +from typing import Optional +import pytest +import torch +import sys + +import triton +import triton.language as tl +from flash_attn.flash_attn_triton_kernel_prefill_amd import MetaData + +DEBUG = False + +def _strides(x: torch.Tensor, *stride_names: str): + if x is None: + return {f"stride_{s}": 0 for i, s in enumerate(stride_names)} + + assert x.ndim == len(stride_names) + return {f"stride_{s}": x.stride(i) for i, s in enumerate(stride_names)} + + +@triton.jit +def _fwd_kernel_splitK( + Q, + K, + V, + sm_scale, + Out_splitK, # [B, H, split_k, Mq, K] + Metadata, # [B, H, 2, split_k, M_ceil] contains [mi, li] + K_new, + V_new, + Cache_seqlens, + Cache_batch_idx, + Alibi_slopes, + stride_qz, + stride_qm, + stride_qg, + stride_qh, + stride_qd, + stride_kz, + stride_kn, + stride_kg, + stride_kh, + stride_kd, + stride_vz, + stride_vn, + stride_vg, + stride_vh, + stride_vd, + stride_osk_zhg, + stride_osk_s, + stride_osk_m, + stride_osk_d, + stride_mzhg, + stride_m2, + stride_ms, + stride_mm, + stride_kn_z, + stride_kn_n, + stride_kn_g, + stride_kn_h, + stride_kn_d, + stride_vn_z, + stride_vn_n, + stride_vn_g, + stride_vn_h, + stride_vn_d, + stride_az, + stride_ah, + Z, + N_CTX_Q, + N_CTX_K, + N_CTX_NEW, + BLOCK_N_PER_SPLIT, + H_q: tl.constexpr, + H_kv: tl.constexpr, + G_q: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + BOUNDS_CHECKS_N: tl.constexpr, + USE_CACHE_SEQLENs: tl.constexpr, + USE_CACHE_BATCH_IDX: tl.constexpr, + NEW_KV: tl.constexpr, + IS_GQA: tl.constexpr, + IS_CAUSAL: tl.constexpr, + USE_ALIBI: tl.constexpr, + PACKED_PER_VAL: tl.constexpr = 1, + N_QUANT_GROUPS: tl.constexpr = 1, +): + """This kernel can accept non-quantized or int4-quantized keys/values. + PACKED_PER_VAL determines the quantization type: + - PACKED_PER_VAL == 1 means no quantization + - PACKED_PER_VAL == 8 means 4-bit quantization (8 packed quantized values inside one int32) + For the quantized case K/V should be int32 tensors. + Quantization can be row-wise (when N_QUANT_GROUPS = 1) or group-wise with N_QUANT_GROUPS = 2, 4, or 8. + Quantization coefficients are stored at the beginning of the row along the last dimension of K/V + So K[B, H, M, :] has a form + [ quant_coef0, quant_coef1, ...| + group0_quant_value0, group0_quant_value1,... | + group1_quant_value0, group1_quant_value1,...] + where each quant_coef is an int32 which should be interpreted as 2 packed float16: scale and offset. + + """ + tl.static_assert( + (PACKED_PER_VAL == 1 and tl.constexpr(K.dtype.element_ty != tl.int32)) + or (PACKED_PER_VAL == 8 and tl.constexpr(K.dtype.element_ty == tl.int32)), + f"Only 4-bit quantization is supported, K/V should have dtype int32 in " + f"the quantized case: {PACKED_PER_VAL=} {tl.constexpr(K.dtype)=} {tl.constexpr(K.dtype.element_ty)=}", + ) + tl.static_assert( + (((N_QUANT_GROUPS == 1 or N_QUANT_GROUPS == 2) or N_QUANT_GROUPS == 4) or N_QUANT_GROUPS == 8), + "Number of quantization groups can be 1 (row-wise quantization), 2, 4, or 8.", + ) + + QUANTIZED: tl.constexpr = PACKED_PER_VAL > 1 + PACKED_D_PER_GROUP: tl.constexpr = BLOCK_DMODEL // PACKED_PER_VAL // N_QUANT_GROUPS + D_PER_GROUP: tl.constexpr = BLOCK_DMODEL // N_QUANT_GROUPS + + start_m = tl.program_id(0) + off_zhg = tl.program_id(1) + off_z = off_zhg // (H_q * G_q) + off_h_q = (off_zhg // G_q) % H_q + off_g_q = off_zhg % G_q + splitk_idx = tl.program_id(2) + + # pick batch index + if USE_CACHE_BATCH_IDX: + cache_batch_idx = tl.load(Cache_batch_idx + off_z) + else: + cache_batch_idx = off_z + + # Load ALiBi slope if enabled + if USE_ALIBI: + a_offset = off_z * stride_az + off_h_q * stride_ah + alibi_slope = tl.load(Alibi_slopes + a_offset) + else: + alibi_slope = None + # print("alibi_slope:", alibi_slope) + + lo = splitk_idx * BLOCK_N_PER_SPLIT + if USE_CACHE_SEQLENs: + cache_seqlen_last_idx = tl.load(Cache_seqlens + off_z) + if NEW_KV: + kv_len = cache_seqlen_last_idx + N_CTX_NEW + else: + kv_len = cache_seqlen_last_idx + else: + kv_len = N_CTX_K + hi = tl.minimum((splitk_idx + 1) * BLOCK_N_PER_SPLIT, kv_len) + + HEAD_RATIO: tl.constexpr = H_q // H_kv + if IS_GQA: + k_head_idx = off_h_q // HEAD_RATIO + v_head_idx = k_head_idx + else: + k_head_idx = off_h_q + v_head_idx = off_h_q + + # calculate base offset + k_base = K + k_head_idx * stride_kh + cache_batch_idx * stride_kz + off_g_q * stride_kg + v_base = V + v_head_idx * stride_vh + cache_batch_idx * stride_vz + off_g_q * stride_vg + + # Copy new Keys and Values into Cache + if NEW_KV: + kn_base = K_new + k_head_idx * stride_kn_h + off_z * stride_kn_z + off_g_q * stride_kn_g + + # Determine the starting position for new data in the cache + if USE_CACHE_SEQLENs: + start_idx = tl.load(Cache_seqlens + off_z) + else: + start_idx = N_CTX_K - N_CTX_NEW + + # Copy new Keys + for i in range(0, N_CTX_NEW, BLOCK_N): + # Load from K_new + k_new_block = tl.load( + kn_base + stride_kn_d * QUANTIZED * N_QUANT_GROUPS + + tl.arange(0, PACKED_D_PER_GROUP)[:, None] * stride_kn_d + + (tl.arange(0, BLOCK_N) + i)[None, :] * stride_kn_n, + mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW), + other=0 + ) + + # Store to K + tl.store( + k_base + stride_kd * QUANTIZED * N_QUANT_GROUPS + + tl.arange(0, PACKED_D_PER_GROUP)[:, None] * stride_kd + + (tl.arange(0, BLOCK_N) + i + start_idx)[None, :] * stride_kn, + k_new_block, + mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) + ) + + # Copy new Values + vn_base = V_new + v_head_idx * stride_vn_h + off_z * stride_vn_z + off_g_q * stride_vn_g + for i in range(0, N_CTX_NEW, BLOCK_N): + # Load from V_new + v_new_block = tl.load( + vn_base + stride_vn_d * QUANTIZED * N_QUANT_GROUPS + + (tl.arange(0, BLOCK_N) + i)[:, None] * stride_vn_n + + tl.arange(0, PACKED_D_PER_GROUP)[None, :] * stride_vn_d, + mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW), + other=0 + ) + + # Store to V + tl.store( + v_base + stride_vd * QUANTIZED * N_QUANT_GROUPS + + (tl.arange(0, BLOCK_N) + i + start_idx)[:, None] * stride_vn + + tl.arange(0, PACKED_D_PER_GROUP)[None, :] * stride_vd, + v_new_block, + mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) + ) + + Q_block_ptr = tl.make_block_ptr( + base=Q + off_h_q * stride_qh + off_z * stride_qz + off_g_q * stride_qg, + shape=(N_CTX_Q, D_PER_GROUP), + strides=(stride_qm, stride_qd), + offsets=(start_m * BLOCK_M, 0), + block_shape=(BLOCK_M, D_PER_GROUP), + order=(1, 0), + ) + + # Additional shift by 1 along the last dimension in the quantized case, since + # the first element along that dim contains packed quantization coefficients. + K_block_ptr = tl.make_block_ptr( + base=k_base + stride_kd * QUANTIZED * N_QUANT_GROUPS, + shape=(PACKED_D_PER_GROUP, hi), + strides=(stride_kd, stride_kn), + offsets=(0, lo), + block_shape=(PACKED_D_PER_GROUP, BLOCK_N), + order=(0, 1), + ) + V_block_ptr = tl.make_block_ptr( + base=v_base + stride_vd * QUANTIZED * N_QUANT_GROUPS, + shape=(hi, PACKED_D_PER_GROUP), + strides=(stride_vn, stride_vd), + offsets=(lo, 0), + block_shape=(BLOCK_N, PACKED_D_PER_GROUP), + order=(1, 0), + ) + + if QUANTIZED: + # Pointers to quantization coefficients + K_scale_shift_block_ptr = tl.make_block_ptr( + base=k_base, + shape=(1, hi), + strides=(stride_kd, stride_kn), + offsets=(0, lo), + block_shape=(1, BLOCK_N), + order=(0, 1), + ) + V_scale_shift_block_ptr = tl.make_block_ptr( + base=v_base, + shape=(hi, 1), + strides=(stride_vn, stride_vd), + offsets=(lo, 0), + block_shape=(BLOCK_N, 1), + order=(1, 0), + ) + else: + K_scale_shift_block_ptr = None + V_scale_shift_block_ptr = None + + # initialize pointer to m and l + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + + acc = tl.zeros([BLOCK_M, D_PER_GROUP], dtype=tl.float32) # noqa: F821 + + # scale sm_scale by log_2(e) and use + # 2^x instead of exp in the loop because CSE and LICM + # don't work as expected with `exp` in the loop + qk_scale = sm_scale * 1.44269504 + # load q: it will stay in SRAM throughout + q = tl.load( # noqa: F821 + tl.advance(Q_block_ptr, (0, 0)), boundary_check=(0, )) + q = (q * qk_scale).to(q.dtype) + # print("q:", q) + + # print("BLOCK_N:", BLOCK_N) + # loop over k, v and update accumulator + for start_n in range(lo, hi, BLOCK_N): + # print("start_n:", start_n) + + k, v = load_dequantize_k_v_group( + K_block_ptr, + V_block_ptr, + K_scale_shift_block_ptr, + V_scale_shift_block_ptr, + BOUNDS_CHECKS_N, + PACKED_PER_VAL, + PACKED_D_PER_GROUP, + Q.dtype.element_ty, + 0, + ) + # print("k:", k) + + # -- compute qk --- + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) # noqa: F821 + # print("qk before:", qk) + + if USE_ALIBI: + row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + col_idx = start_n + tl.arange(0, BLOCK_N) + + # Compute relative positions + relative_pos = row_idx[:, None] + kv_len - (N_CTX_Q + col_idx[None, :]) + relative_pos = tl.abs(relative_pos) + # print("relative_pos:", relative_pos) + + # Compute ALiBi bias + alibi_bias = -1 * alibi_slope * relative_pos + # print("alibi_bias:", alibi_bias) + qk += (alibi_bias * 1.44269504) + + # Apply causal mask if IS_CAUSAL is True + if IS_CAUSAL: + row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + # print("row_idx:", row_idx) + col_idx = start_n + tl.arange(0, BLOCK_N) + # print("col_idx:", col_idx) + + # create a N_CTX_Q x kv_len causal mask + col_offset = N_CTX_Q - kv_len + causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :]) + + # Apply the mask + qk = tl.where(causal_mask, qk, float("-inf")) + # print("qk after causal:", qk) + + # TODO: This is slow, and only needed at the last iteration. + # Maybe we can unroll the last iteration instead? + if BOUNDS_CHECKS_N: + qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf")) + # print("qk after BOUNDS_CHECKS_N:", qk) + + # -- compute scaling constant --- + # print("m_i:", m_i) + m_i_new = tl.maximum(m_i, tl.max(qk, 1)) + # print("m_i_new:", m_i_new) + if IS_CAUSAL: + alpha = tl.math.exp2(tl.where(m_i > float("-inf"), m_i - m_i_new, float("-inf"))) + else: + alpha = tl.math.exp2(m_i - m_i_new) + + # print("alpha:", alpha) + # print("before qk - m_i_new:", qk) + # cause of nan because subtracting infs + if IS_CAUSAL: + qk = tl.where(qk > float("-inf"), qk - m_i_new[:, None], float("-inf")) + else: + qk = qk - m_i_new[:, None] + + p = tl.math.exp2(qk) + # print("p:", p) + + # -- update m_i and l_i -- + l_i = l_i * alpha + tl.sum(p, 1) + m_i = m_i_new + p = p.to(Q.dtype.element_ty) + + # -- scale and update acc -- + acc *= alpha[:, None] + acc += tl.dot(p, v) + # print("acc:", acc) + + # update pointers + K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + if PACKED_PER_VAL > 1: + K_scale_shift_block_ptr = tl.advance(K_scale_shift_block_ptr, (0, BLOCK_N)) + V_scale_shift_block_ptr = tl.advance(V_scale_shift_block_ptr, (BLOCK_N, 0)) + + # write back O + O_block_ptr = tl.make_block_ptr( + base=Out_splitK + off_zhg * stride_osk_zhg + splitk_idx * stride_osk_s, + shape=(N_CTX_Q, D_PER_GROUP), + strides=(stride_osk_m, 1), + offsets=(start_m * BLOCK_M, 0), + block_shape=(BLOCK_M, D_PER_GROUP), + order=(1, 0), + ) + tl.store( + tl.advance(O_block_ptr, (0, 0)), + acc, + boundary_check=(0, ), + ) + # Write metadata for split-K reduction + Metadata_ptr = (Metadata + off_zhg * stride_mzhg + splitk_idx * stride_ms + start_m * BLOCK_M + + tl.arange(0, BLOCK_M)) + tl.store(Metadata_ptr, m_i) + tl.store(Metadata_ptr + stride_m2, l_i) + + +@triton.jit +def load_dequantize_k_v_group( + K_block_ptr, + V_block_ptr, + K_scale_shift_block_ptr, + V_scale_shift_block_ptr, + BOUNDS_CHECKS_N: tl.constexpr, + PACKED_PER_VAL: tl.constexpr, + PACKED_D_PER_GROUP: tl.constexpr, + dtype: tl.constexpr, + group_id: tl.constexpr, +): + #Load K/V for a given block. In case of int4-quantized K/V, + # dequantize them after loading. If quantization is group-wise, + # use group_id to advance the pointers to the current group. + + # Advance to the current quantization group + K_block_ptr = tl.advance(K_block_ptr, (PACKED_D_PER_GROUP * group_id, 0)) + V_block_ptr = tl.advance(V_block_ptr, (0, PACKED_D_PER_GROUP * group_id)) + + # -- load k, v -- + k = tl.load(K_block_ptr, boundary_check=(1, ) if BOUNDS_CHECKS_N else ()) + v = tl.load(V_block_ptr, boundary_check=(0, ) if BOUNDS_CHECKS_N else ()) + + if PACKED_PER_VAL > 1: + # K/V are quantized, load quantization coefficients and dequantize + K_scale_shift_block_ptr = tl.advance(K_scale_shift_block_ptr, (group_id, 0)) + V_scale_shift_block_ptr = tl.advance(V_scale_shift_block_ptr, (0, group_id)) + + k_scale_shift = tl.load(K_scale_shift_block_ptr, boundary_check=(1, ) if BOUNDS_CHECKS_N else ()) + v_scale_shift = tl.load(V_scale_shift_block_ptr, boundary_check=(0, ) if BOUNDS_CHECKS_N else ()) + + k_scale, k_shift = cast_uint32_to_half2(k_scale_shift) + v_scale, v_shift = cast_uint32_to_half2(v_scale_shift) + v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL).to(dtype) + k_t = dequantize( + tl.trans(k), + tl.trans(k_scale), + tl.trans(k_shift), + PACKED_PER_VAL, + ).to(dtype) + k = tl.trans(k_t) + return k, v + + +@triton.jit +def cast_uint32_to_half2(scale_shift): + # Extract two float16 packed into one int32 + scale = scale_shift & 0xFFFF + shift = scale_shift >> 16 + scale = scale.to(tl.uint16).to(tl.float16, bitcast=True) + shift = shift.to(tl.uint16).to(tl.float16, bitcast=True) + return scale, shift + + +@triton.jit +def dequantize( + x_, + scale, + shift, + PACKED_PER_VAL: tl.constexpr = 8, +): + # PACKED_PER_VAL is the number of values packed into + # each element x_. For example, for int4 quantization + #and x_ of type int32, PACKED_PER_VAL is 8. + + BLOCK_N: tl.constexpr = x_.shape[0] + BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1] + offsets = tl.arange(0, PACKED_PER_VAL) * 4 + quant_offset = (x_[:, None, :] >> offsets[None, :, None]) # (BLOCK_N, PACKED_PER_VAL, D // PACKED_PER_VAL) + + quant_offset = tl.view(quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)) + # Trick - instead of converting int4 to float16 we view it as float16 + # and then multiply by 32768 * 512 == 2**24 + quant_offset = (quant_offset & 0xF).to(tl.uint16).to(tl.float16, bitcast=True) + quant_offset = (quant_offset * 32768.0).to(tl.float16) + scale_512 = scale * 512 + + dequant = quant_offset * scale_512 + shift + return dequant + + +@triton.jit +def _splitK_reduce( + Out_splitK, # [B, H, split_k, Mq, K] + Metadata, # [B, H, 2, split_k, M_ceil] contains [mi, li] + Out, # [B, H, M, K] + LSE, # [B, H, M] + stride_osk_zhg, + stride_osk_s, + stride_osk_m, + stride_osk_k, + stride_mzhg, + stride_m2, + stride_ms, + stride_mm, + stride_oz, + stride_oh, + stride_og, + stride_om, + stride_ok, + stride_lse_zhg, + stride_lse_m, + M_ceil: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + H: tl.constexpr, + G: tl.constexpr, + split_k: tl.constexpr, + splitK_pow2: tl.constexpr, + use_mask: tl.constexpr, + IS_CAUSAL: tl.constexpr, +): + off_zhg = tl.program_id(0) + off_z = off_zhg // (H * G) + off_h = (off_zhg // G) % H + off_g = off_zhg % G + off_m = tl.program_id(1) + off_k = tl.program_id(2) + + # read chunk + spk_idx = tl.arange(0, splitK_pow2) + kidx = tl.arange(0, BLOCK_SIZE) + + Metadata_ptr = (Metadata + stride_mzhg * off_zhg + spk_idx * stride_ms + off_m * stride_mm) + + o_ptr = (Out_splitK + off_zhg * stride_osk_zhg + stride_osk_m * off_m + off_k * BLOCK_SIZE + + stride_osk_s * spk_idx[:, None] + kidx[None, :] * stride_osk_k) + + # read max values of each splitK + if use_mask: + spk_mask = spk_idx < split_k + l_m = tl.load(Metadata_ptr, mask=spk_mask, other=float("-inf")) + l_sum = tl.load(Metadata_ptr + stride_m2, mask=spk_mask, other=0.0) + acc = tl.load(o_ptr, mask=spk_mask[:, None], other=0.0) + else: + l_m = tl.load(Metadata_ptr) + l_sum = tl.load(Metadata_ptr + stride_m2) + acc = tl.load(o_ptr) + + # print("l_m:", l_m) + # print("l_sum:", l_sum) + # print("acc:", acc) + + g_m = tl.max(l_m, axis=0) + + # print("g_m:", g_m) + if IS_CAUSAL: + l_m_offset = l_m - g_m + alpha = tl.where(l_m_offset > float("-inf"), tl.math.exp2(l_m_offset), 0.0) + else: + alpha = tl.math.exp2(l_m - g_m) + # print("alpha:", alpha) + + # read sum + l_sum *= alpha + g_sum = tl.sum(l_sum, axis=0) + acc = acc * alpha[:, None] + + if IS_CAUSAL: + # Avoid division by zero + g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0) + acc_out = tl.sum(acc, axis=0) / g_sum_safe + else: + acc_out = tl.sum(acc, axis=0) / g_sum + + # Store output + Out_ptr = (Out + stride_oz * off_z + stride_oh * off_h + stride_og * off_g + stride_om * off_m + + off_k * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) + tl.store(Out_ptr, acc_out) + + # Store lse + l_ptrs = LSE + off_zhg * stride_lse_zhg + off_m + if IS_CAUSAL: + lse = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m) + tl.store(l_ptrs, lse) + else: + tl.store(l_ptrs, (g_m + tl.math.log2(g_sum)) / 1.44269504) + + +def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor: + # Scale and shift are such that quantization linearly maps + # int4 values range [0..15] to input values range min(k)..max(k) + # individually for every row + k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups) + max_vals = torch.max(k, dim=-1, keepdim=True).values + min_vals = torch.min(k, dim=-1, keepdim=True).values + scale_k: torch.Tensor = (max_vals - min_vals) / 15 + + shift_k = torch.min(k, dim=-1, keepdim=True).values + scale_k = scale_k.to(torch.float16) + shift_k = shift_k.to(torch.float16) + + in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5 + in_bytes = in_bytes.to(torch.uint8) + in_int4 = in_bytes & 0xF + in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4) + scale_shift = torch.concat([scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1) + k_quant = torch.concat( + [ + scale_shift.flatten(start_dim=-2), + in_int4_packed.flatten(start_dim=-2), + ], + dim=-1, + ).view(torch.int16) + return k_quant + + +def dequantize_kv_fp16(quant_k: torch.Tensor, num_groups: int = 1) -> torch.Tensor: + k_i16 = quant_k.view(torch.int16) + k_ui8 = k_i16.view(torch.uint8) + + ss_size = num_groups * 4 + scale_shift_ui8 = k_ui8[..., 0:ss_size] + scale_shift_ui8 = scale_shift_ui8.reshape(*scale_shift_ui8.shape[:-1], num_groups, 4) + scale = scale_shift_ui8[..., 0:2].view(torch.float16) + shift = scale_shift_ui8[..., 2:4].view(torch.float16) + + kv_ui8 = k_ui8[..., ss_size:] + k_ui8 = kv_ui8.reshape(*kv_ui8.shape[:-1], num_groups, -1) + k1_i4 = k_ui8 & 0xF + k2_i4 = (k_ui8 & 0xF0) >> 4 + k_shape = k1_i4.shape + k1_f16 = k1_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape) + k2_f16 = k2_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape) + + out = torch.empty((*k1_f16.shape[:-1], k1_f16.shape[-1] * 2), dtype=torch.float16, device=quant_k.device) + out[..., ::2] = k1_f16 + out[..., 1::2] = k2_f16 + out = out.reshape(*k_shape[:-2], -1) + + return out + + +def get_split_k(B: int, G: int, H: int, Mk: int) -> int: + """Heuristic for the number of splits""" + bh = max(B * H, 1) # NOTE: Handle B*h=0 case + split_k = max(Mk, 1024) // bh + max_chunk_size = 64 + while split_k > 0 and Mk / split_k < max_chunk_size: + split_k = split_k // 2 + while B * H * G * split_k >= 1024: + split_k = split_k // 2 + split_k = min(split_k, 512) + split_k = max(split_k, 1) + return split_k + +class _attention(torch.autograd.Function): + + OPERATOR = _fwd_kernel_splitK + SUPPORTED_DEVICES = {"cuda"} + CUDA_MINIMUM_COMPUTE_CAPABILITY = (8, 0) + SUPPORTED_DTYPES = { + torch.half, + torch.bfloat16, + } + SUPPORTED_MAX_K = 128 + SUPPORTS_DROPOUT = False + SUPPORTS_CUSTOM_SCALE = True + SUPPORTS_BMGHK = True + NAME = "triton_splitKF" + + @staticmethod + def forward(cls, q, k, v, input_metadata): + if DEBUG: + print() + print("attention_decode.forward") + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("input_metadata:", input_metadata) + + original_layout = input_metadata.layout + + # kernels expects "bsghd" + if input_metadata.layout == "bshd": + q=q.unsqueeze(2) + k=k.unsqueeze(2) + v=v.unsqueeze(2) + + if input_metadata.new_kv: + input_metadata.k_new = input_metadata.k_new.unsqueeze(2) + input_metadata.v_new = input_metadata.v_new.unsqueeze(2) + + input_metadata.layout = "bsghd" + elif input_metadata.layout == "bhsd": + q=q.permute(0, 2, 1, 3).unsqueeze(2) + k=k.permute(0, 2, 1, 3).unsqueeze(2) + v=v.permute(0, 2, 1, 3).unsqueeze(2) + if input_metadata.new_kv: + input_metadata.k_new = input_metadata.k_new.permute(0, 2, 1, 3).unsqueeze(2) + input_metadata.v_new = input_metadata.v_new.permute(0, 2, 1, 3).unsqueeze(2) + + + input_metadata.layout = "bsghd" + elif input_metadata.layout == "bsghd": + pass + elif input_metadata.layout is None: + raise ValueError("Layout not given") + + assert input_metadata.layout == "bsghd" + + # get dims + batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_q = q.shape + _, seqlen_k, n_group_k, heads_per_group_k, dim_k = k.shape + _, seqlen_v, n_group_v, heads_per_group_v, dim_v = v.shape + + print(f"batch_size = {batch_size}, seqlen_q = {seqlen_q}, group_q={n_group_q} heads_per_group_q = {heads_per_group_q}, dim_q = {dim_q}") + print(f"batch_size = {batch_size}, seqlen_k = {seqlen_k}, group_k={n_group_k} heads_per_group_k = {heads_per_group_k}, dim_k = {dim_k}") + + # Handle MQA/GQA case + if heads_per_group_q > heads_per_group_k: + input_metadata.is_gqa = True + elif heads_per_group_q < heads_per_group_k: + raise ValueError("heads_per_group_q < heads_per_group_k") + else: + input_metadata.is_gqa = False + + print("input_metadata.is_gqa:", input_metadata.is_gqa) + print("After MQA/GQA check") + print(f"batch_size = {batch_size}, seqlen_q = {seqlen_q}, group_q={n_group_q} heads_per_group_q = {heads_per_group_q}, dim_q = {dim_q}") + print(f"batch_size = {batch_size}, seqlen_k = {seqlen_k}, group_k={n_group_k} heads_per_group_k = {heads_per_group_k}, dim_k = {dim_k}") + + # context + cls.SPLIT_K: Optional[int] = None + cls.BLOCK_M = 16 + cls.BLOCK_N = 64 + + cls.NUM_QUANT_GROUPS = 1 # Default quantization is row-wise + + # attn_bias = inp.attn_bias + if input_metadata.cache_seqlens is not None: + cache_seqlens = input_metadata.cache_seqlens + else: + cache_seqlens = None + + # Transpose in the case of MQA/GQA + mqa_swap_seqlen_head = False + # if heads_per_group_k > 1 and k.stride(3) == 0 and v.stride(3) == 0: + # mqa_swap_seqlen_head = True + # assert seqlen_q == 1 + # q = q.transpose(1, 3) + # k = k[:, :, :, :1] + # v = v[:, :, :, :1] + print("mqa_swap_seqlen_head:", mqa_swap_seqlen_head) + # assert mqa_swap_seqlen_head == False + + # Update dim_k if Quantized + PACKED_PER_VAL = 1 + if k.dtype == torch.int32: + # Quantized K/V + PACKED_PER_VAL = 8 + dim_k = (dim_k - cls.NUM_QUANT_GROUPS) * 8 + + assert dim_k == dim_q, f"Keys have head dim {dim_k} but queries have head dim {dim_q}" + # print(f"batch_size = {batch_size}, seqlen_q = {seqlen_q}, seqlen_k = {seqlen_k}, heads_per_group_q = {heads_per_group_q}, heads_per_group_k = {heads_per_group_k}, dim_q = {dim_q}, dim_k = {dim_k}") + + BLOCK_M = cls.BLOCK_M + BLOCK_N = cls.BLOCK_N + if cls.SPLIT_K is not None: + split_k = cls.SPLIT_K + else: + # Use heuristics + split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, seqlen_k) # NOTE: should the split think about seqlens? + if DEBUG: + print("split_k:", split_k) + + seqlen_q_ceil = (seqlen_q + BLOCK_M - 1) // BLOCK_M * BLOCK_M + out_splitk = torch.empty([batch_size * n_group_q * heads_per_group_q, split_k, seqlen_q_ceil, dim_q], dtype=torch.float32, device=q.device) + metadata = torch.empty([batch_size * n_group_q * heads_per_group_q, 2, split_k, seqlen_q_ceil], dtype=torch.float32, device=q.device) + lse = torch.empty((batch_size * n_group_q * heads_per_group_q, seqlen_q), device=q.device, dtype=torch.float32) + grid = (triton.cdiv(seqlen_q, BLOCK_M), batch_size * n_group_q * heads_per_group_q, split_k) + + num_warps = 1 + split_size = (seqlen_k + split_k - 1) // split_k + use_cache_seqlens = cache_seqlens is not None + + print(f"batch_size = {batch_size}, group_q = {n_group_q}, heads_per_group_q = {heads_per_group_q}, split_k = {split_k}, seqlen_q_ceil = {seqlen_q_ceil}, dim_q = {dim_q}, num_of_wgs = {n_group_q * n_group_q * heads_per_group_q * split_k}") + + if DEBUG: + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("sm_scale:", input_metadata.sm_scale) + print("o_splitk:", out_splitk, out_splitk.shape) + print("metadata:", metadata, metadata.shape) + print("cache_seqlens:", cache_seqlens) + print("lse:", lse) + print("grid:", grid) + print("split_size:", split_size) + print("use_cache_seqlens:", use_cache_seqlens) + + + _fwd_kernel_splitK[grid]( + Q=q, + K=k, + V=v, + sm_scale=input_metadata.sm_scale, + Out_splitK=out_splitk, + Metadata=metadata, + K_new = input_metadata.k_new, + V_new = input_metadata.v_new, + Cache_seqlens=cache_seqlens, + Cache_batch_idx=input_metadata.cache_batch_idx, + Alibi_slopes=input_metadata.alibi_slopes, + **_strides(q, "qz", "qm", "qg", "qh", "qd"), + **_strides(k, "kz", "kn", "kg", "kh", "kd"), + **_strides(v, "vz", "vn", "vg", "vh", "vd"), + **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_d"), + **_strides(metadata, "mzhg", "m2", "ms", "mm"), + **_strides(input_metadata.k_new, "kn_z", "kn_n", "kn_g", "kn_h", "kn_d"), + **_strides(input_metadata.v_new, "vn_z", "vn_n", "vn_g", "vn_h", "vn_d"), + **_strides(input_metadata.alibi_slopes, "az", "ah"), + Z=batch_size, + H_q=heads_per_group_q, + H_kv=heads_per_group_k, + G_q=n_group_q, + N_CTX_Q=seqlen_q, + N_CTX_K=seqlen_k, + N_CTX_NEW=input_metadata.k_new.shape[1] if input_metadata.new_kv else None, + BLOCK_N_PER_SPLIT=split_size, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=dim_k, + BOUNDS_CHECKS_N=(split_size % BLOCK_N) > 0 or use_cache_seqlens, + USE_CACHE_SEQLENs=use_cache_seqlens, + USE_CACHE_BATCH_IDX= input_metadata.cache_batch_idx is not None, + NEW_KV=input_metadata.new_kv, + IS_GQA=input_metadata.is_gqa, + IS_CAUSAL=input_metadata.causal, + USE_ALIBI=False if input_metadata.alibi_slopes is None else True, + num_warps=num_warps, + num_stages=1, + PACKED_PER_VAL=PACKED_PER_VAL, + N_QUANT_GROUPS=cls.NUM_QUANT_GROUPS if PACKED_PER_VAL > 1 else 1, + ) + + if mqa_swap_seqlen_head: + out = torch.empty((batch_size, heads_per_group_q, n_group_q, seqlen_q, dim_q), device=q.device, dtype=q.dtype).transpose(1, 3) + else: + out = torch.empty((batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_q), device=q.device, dtype=q.dtype) + + # Merge together + splitK_pow2 = triton.next_power_of_2(split_k) + use_mask = splitK_pow2 > split_k + if batch_size * n_group_q * heads_per_group_q * seqlen_q >= 512: + k_block_num = 1 + else: + k_block_num = 2 + assert out.shape[-1] % k_block_num == 0 + k_block_size = out.shape[-1] // k_block_num + grid = (batch_size * n_group_q * heads_per_group_q, seqlen_q, k_block_num) + print("grid:", grid) + + + if DEBUG: + print("Before _splitK_reduce call") + print("out_splitk:", out_splitk, out_splitk.shape) + print("metadata:", metadata, metadata.shape) + print("out:", out) + print("lse:", lse) + print("use_mask:", use_mask) + + + _splitK_reduce[grid]( + out_splitk, + metadata, + out, + lse, + **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_k"), + **_strides(metadata, "mzhg", "m2", "ms", "mm"), + **_strides(out, "oz", "om", "og", "oh", "ok"), + **_strides(lse, "lse_zhg", "lse_m"), + M_ceil=seqlen_q_ceil, + BLOCK_SIZE=k_block_size, + G=n_group_q, + H=heads_per_group_q, + # TODO: Tune num_warps + split_k=split_k, + splitK_pow2=splitK_pow2, + use_mask=use_mask, + IS_CAUSAL=input_metadata.causal, + num_warps=4) + + lse = lse.reshape([batch_size, n_group_q, heads_per_group_q, seqlen_q]) + if mqa_swap_seqlen_head: + # H/M dimensions have been swapped + out = out.transpose(1, 3) + lse = lse.transpose(2, 3) + if q.ndim == 4: + # BMGHK -> BMHK + assert n_group_q == 1 + out = out[:, :, 0] + lse = lse[:, 0] + if seqlen_k == 0: + out.zero_() + if mqa_swap_seqlen_head: + out = out.reshape(batch_size, -1, seqlen_q * n_group_q, dim_q).transpose(1, 2).contiguous() + else: + out = out.reshape(batch_size, heads_per_group_q * n_group_q, -1, dim_q).contiguous() + + # output is batch_size, heads_per_group_q * group_q, seqlen_q, dim_q + if original_layout == "bshd": + # out=out.transpose(1, 2).contiguous() # this screws up heads and data. + # the data is laid out properly. Just need to reshape dims + out = out.reshape(batch_size, seqlen_q, -1, dim_q) + + return out + + +attention_decode = _attention.apply + + +def get_input_shapes(): + cases = [(max(1, 2**(16 - i)), 1, 2**i, 16, 1, 128) + for i in range(8, 18)] + [(max(1, 2**(16 - i)), 1, 2**i, 16, 2, 128) for i in range(8, 18)] + + return cases + + +@pytest.mark.parametrize('batch_size, seqlen_q, seqlen_k, group_q, group_k, dim', get_input_shapes()) +def test_op_fwd(batch_size, seqlen_q, seqlen_k, group_q, group_k, dim, dtype=torch.bfloat16): + print() + print(f"batch_size = {batch_size}, seqlen_q = {seqlen_q}, seqlen_k = {seqlen_k}, group_q = {group_q}, group_k = {group_k}, dim = {dim}") + torch.manual_seed(20) + query_group_head_size = (group_q + group_k - 1) // group_k + q = (torch.empty((batch_size, seqlen_q, group_k, query_group_head_size, dim), dtype=dtype, + device="cuda").normal_(mean=0., std=0.5).requires_grad_()) + k = (torch.empty((batch_size, seqlen_k, group_k, 1, dim), dtype=dtype, + device="cuda").normal_(mean=0., + std=0.5).requires_grad_()).expand(-1, -1, -1, query_group_head_size, -1) + v = (torch.empty((batch_size, seqlen_k, group_k, 1, dim), dtype=dtype, + device="cuda").normal_(mean=0., + std=0.5).requires_grad_()).expand(-1, -1, -1, query_group_head_size, -1) + scale = 1 / dim**0.5 + + print("q:", q.shape) + print("k:", k.shape) + print("v:", v.shape) + input_metadata = MetaData(sm_scale=scale) + input_metadata.layout = "bsghd" + tri_out = attention_decode(q, k, v, input_metadata) + print("tri_out:", tri_out.shape) + print() + + q = q.reshape([batch_size, seqlen_q, -1, dim]).permute(0, 2, 1, 3) + k = k.reshape([batch_size, seqlen_k, -1, dim]).permute(0, 2, 1, 3) + v = v.reshape([batch_size, seqlen_k, -1, dim]).permute(0, 2, 1, 3) + print("q_ref:", q.shape) + print("k_ref:", k.shape) + print("v_ref:", v.shape) + attn = (q @ k.transpose(-1, -2) * scale).softmax(-1) + print("attn:", attn.shape) + ref_out = attn @ v + print("ref_out:", ref_out.shape) + + # compare + torch.testing.assert_close(ref_out, tri_out, atol=1e-3, rtol=0) + + +@pytest.mark.parametrize('B, Mq, Mkv, Hq, Hkv, K', get_input_shapes()) +def test_op_fwd_int4_kv(B, Mq, Mkv, Hq, Hkv, K, dtype=torch.float16): + torch.manual_seed(2) + q = (torch.empty((B, Mq, Hkv, (Hq + Hkv - 1) // Hkv, K), dtype=dtype, + device="cuda").normal_(mean=1.0, std=0.5).requires_grad_()) + k = (torch.empty((B, Mkv, Hkv, 1, K), dtype=dtype, + device="cuda").normal_(mean=1.0, + std=0.5).requires_grad_()).expand(-1, -1, -1, (Hq + Hkv - 1) // Hkv, -1) + v = (torch.empty((B, Mkv, Hkv, 1, K), dtype=dtype, + device="cuda").normal_(mean=1.0, + std=0.5).requires_grad_()).expand(-1, -1, -1, (Hq + Hkv - 1) // Hkv, -1) + + num_groups = 1 + quant_k = (quantize_kv_int4(k, num_groups=num_groups).contiguous().view(torch.int32)) + quant_v = (quantize_kv_int4(v, num_groups=num_groups).contiguous().view(torch.int32)) + scale = 1 / K**0.5 + input_metadata = MetaData(sm_scale=scale) + input_metadata.layout = "bsghd" + tri_out = attention_decode(q, quant_k, quant_v, input_metadata) + + q = q.reshape([B, Mq, -1, K]).permute(0, 2, 1, 3) + k = k.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3) + v = v.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3) + attn = (q @ k.transpose(-1, -2) * scale).softmax(-1) + ref_out = attn @ v + # compare + torch.testing.assert_close(ref_out, tri_out, atol=2.1e-2, rtol=0) + + # since quantization introduces rounding error, use the + # dequantized kv as inputs to the ref implementation to reduce + # the tolerance to 1e-3 + dqk = dequantize_kv_fp16(quant_k, num_groups=num_groups) + dqv = dequantize_kv_fp16(quant_v, num_groups=num_groups) + dqk = dqk.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3) + dqv = dqv.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3) + dq_attn = (q @ dqk.transpose(-1, -2) * scale).softmax(-1) + dq_ref_out = dq_attn @ dqv + torch.testing.assert_close(dq_ref_out, tri_out, atol=1e-3, rtol=0) + + +def test_quantization(): + a = torch.randn((2, 4, 32), dtype=torch.float16, device='cuda') + qa = quantize_kv_int4(a, num_groups=4) + dqa = dequantize_kv_fp16(qa, num_groups=4) + torch.testing.assert_close(a, dqa, atol=1.5e-1, rtol=1e-1) + + +try: + FLASH_VER = 2 +except BaseException: + try: + FLASH_VER = 1 + except BaseException: + FLASH_VER = None +HAS_FLASH = FLASH_VER is not None + +configs = [] +for mode in ['fwd']: + # for D_HEAD in [128]: + for causal in [False]: + configs.append( + triton.testing.Benchmark( + x_names=['B', 'Mq', 'Mkv', 'Hq', 'Hkv', 'K'], x_vals=get_input_shapes(), line_arg='provider', + line_vals=['triton'] + (['flash'] if HAS_FLASH else []), + line_names=['Triton'] + ([f'Flash-{FLASH_VER}'] if HAS_FLASH else []), styles=[('red', '-'), + ('blue', '-')], + ylabel='ms', plot_name=f'fused-attention-d{128}-{mode}-causal={causal}', args={ + # 'D_HEAD': D_HEAD, + 'dtype': torch.float16, 'mode': mode, 'causal': causal + })) + + +@triton.testing.perf_report(configs) +def bench_flash_attention(B, Mq, Mkv, Hq, Hkv, K, causal, mode, provider, dtype=torch.float16, device="cuda"): + assert mode in ['fwd', 'bwd'] + warmup = 100 + rep = 400 + ms = 0 + if provider == "triton": + q = torch.randn([B, Mq, Hkv, Hq // Hkv, K], device="cuda", dtype=dtype, requires_grad=False) + k = torch.randn([B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, + requires_grad=False).expand(-1, -1, -1, Hq // Hkv, -1) + v = torch.randn([B, Mkv, Hkv, 1, K], device="cuda", dtype=dtype, + requires_grad=False).expand(-1, -1, -1, Hq // Hkv, -1) + + sm_scale = 1.3 + input_metadata = MetaData(sm_scale=sm_scale) + input_metadata.layout = "bsghd" + fn = lambda: attention_decode(q, k, v, input_metadata) + ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep) + + # flops_per_matmul = 2 * B * Hq * (Mq * K * Mkv + Mq * Mkv * K) + # total_flops = 2 * flops_per_matmul + # totalBytes = ((B * Mkv * Hkv * K * 2) + (B * Mq * Hq * K) + (B * Mq * Hq * K)) * 2 + + # return totalBytes / ms * 1e-9 + return ms * 1000 + + +def main(): + bench_flash_attention.run(save_path='.', print_data=True) + + +if __name__ == '__main__': + sys.exit(main()) \ No newline at end of file diff --git a/flash_attn/flash_attn_triton_kernel_amd.py b/flash_attn/flash_attn_triton_kernel_prefill_amd.py similarity index 98% rename from flash_attn/flash_attn_triton_kernel_amd.py rename to flash_attn/flash_attn_triton_kernel_prefill_amd.py index a13fc6c3301..1b1bc268331 100644 --- a/flash_attn/flash_attn_triton_kernel_amd.py +++ b/flash_attn/flash_attn_triton_kernel_prefill_amd.py @@ -28,6 +28,8 @@ import triton import triton.language as tl +DEBUG = False + class MetaData(): cu_seqlens_q = None @@ -40,10 +42,17 @@ class MetaData(): num_contexts = 0 varlen = False layout = None + cache_seqlens = None + cache_batch_idx = None + new_kv = False + seqlen_new = None + k_new = None + v_new = None dropout_p, return_encoded_softmax = 0.0, False def __repr__(self) -> str: return (f"MetaData(\n" + f" sm_scale={self.sm_scale},\n" f" cu_seqlens_q={self.cu_seqlens_q},\n" f" cu_seqlens_k={self.cu_seqlens_k},\n" f" max_seqlens_q={self.max_seqlens_q},\n" @@ -54,6 +63,12 @@ def __repr__(self) -> str: f" num_contexts={self.num_contexts},\n" f" varlen={self.varlen},\n" f" layout={self.layout},\n" + f" cache_seqlens={self.cache_seqlens},\n" + f" cache_batch_idx={self.cache_batch_idx},\n" + f" new_kv={self.new_kv},\n" + f" seqlen_new={self.seqlen_new},\n" + f" k_new={self.k_new},\n" + f" v_new={self.v_new},\n" f" dropout_p={self.dropout_p},\n" f" return_encoded_softmax={self.return_encoded_softmax}\n" f")") @@ -259,12 +274,14 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri size_n = start_n + OFFS_N[None, :] mask = size_n < boundary_m[:, None] qk = tl.where(mask, qk, float("-inf")) + + # -- compute qk ---- + qk += tl.dot(q, k) + if IS_CAUSAL: causal_boundary = start_n + offs_n_causal causal_mask = OFFS_M[:, None] >= causal_boundary[None, :] qk = tl.where(causal_mask, qk, float("-inf")) - # -- compute qk ---- - qk += tl.dot(q, k) if bias_ptrs is not None: bias_offs_n = start_n + tl.arange(0, BLOCK_N) if MASK_STEPS else None bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k) @@ -842,10 +859,6 @@ def _attn_bwd(Q, K, V, sm_scale, alibi_slopes, DO, DQ, DK, DV, M, D, dq *= LN2 tl.store(DQ_block_ptr, dq.to(q.dtype)) - -empty = torch.empty(128, device="cuda") - - def get_shape_from_layout(q, k, metadata): if metadata.layout == 'thd': nheads_q, nheads_k = q.shape[1], k.shape[1] @@ -888,6 +901,15 @@ class _attention(torch.autograd.Function): @staticmethod def forward(ctx, q, k, v, o, metadata): + if DEBUG: + print() + print("_attention.forward") + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("o:", o, o.shape) + print("metadata:", metadata) + # NOTE: a large bias tensor leads to overflow during pointer arithmetic if (metadata.bias is not None): assert (metadata.bias.numel() < 2**31) @@ -1031,7 +1053,7 @@ def backward(ctx, do, _): return dq, dk, dv, None, None -attention = _attention.apply +attention_prefill = _attention.apply def input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout): @@ -1126,7 +1148,7 @@ def test_op_fwd(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_alibi, layout, o = torch.empty_like(q) # triton implementation - tri_out, _ = attention(q, k, v, o, input_metadata) + tri_out, _ = attention_prefill(q, k, v, o, input_metadata) # Transpose here if layout is bshd so we have same reference code for all layouts if layout == 'bshd': @@ -1197,7 +1219,7 @@ def test_op_fwd_bias(Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_bias, dtype=tor o = torch.empty_like(q) # triton implementation - tri_out, _ = attention(q, k, v, o, input_metadata) + tri_out, _ = attention_prefill(q, k, v, o, input_metadata) # reference implementation:171 scores = torch.einsum('bhqd,bhkd->bhqk', q, k).float() * sm_scale @@ -1236,7 +1258,7 @@ def test_op_varlen_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16): scores = torch.einsum('qhd,khd->qhk', q[start_q:end_q], k[start_k:end_k]).float() p = torch.softmax(scores * input_metadata.sm_scale, dim=-1).half() ref_out[start_q:end_q] = torch.einsum('qhk,khd->qhd', p, v[start_k:end_k]) - attention(q, k, v, tri_out, input_metadata) + attention_prefill(q, k, v, tri_out, input_metadata) torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2) @@ -1264,7 +1286,7 @@ def test_op_varlen_mqa_fwd(Z, HQ, HK, N_CTX, D_HEAD, causal, dtype=torch.float16 scores = torch.einsum('qhd,khd->qhk', q[start_q:end_q], k_curr).float() p = torch.softmax(scores * input_metadata.sm_scale, dim=-1).half() ref_out[start_q:end_q] = torch.einsum('qhk,khd->qhd', p, v_curr) - attention(q, k, v, tri_out, input_metadata) + attention_prefill(q, k, v, tri_out, input_metadata) torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=1e-2) @@ -1341,7 +1363,7 @@ def test_op_bwd(Z, H, N_CTX, D_HEAD, qseqlen_not_equal_kseqlen, causal, torch_sd ref_dq, q.grad = q.grad.clone(), None # # triton implementation - tri_out, _ = attention(q, k, v, o, input_metadata) + tri_out, _ = attention_prefill(q, k, v, o, input_metadata) tri_out.backward(dout) tri_dv, v.grad = v.grad.clone(), None tri_dk, k.grad = k.grad.clone(), None @@ -1473,7 +1495,7 @@ def bench_flash_attention(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, causal if causal: input_metadata.need_causal() o = torch.empty_like(q) - fn = lambda: attention(q, k, v, o, input_metadata) + fn = lambda: attention_prefill(q, k, v, o, input_metadata) if mode == 'bwd': o, _ = fn() do = torch.randn_like(o) diff --git a/tests/test_flash_attn.py b/tests/test_flash_attn.py index 28e28a36c85..7051b1c6e1c 100644 --- a/tests/test_flash_attn.py +++ b/tests/test_flash_attn.py @@ -17,6 +17,8 @@ from flash_attn.flash_attn_interface import _get_block_size_n from flash_attn.layers.rotary import apply_rotary_emb +DEBUG = False + MAX_HEADDIM_SM8x = 192 @@ -48,6 +50,8 @@ def attn_bias_from_alibi_slopes( else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") ) relative_pos = torch.abs(row_idx + sk - sq - col_idx) + if DEBUG: + print("relative_pos:", relative_pos) return -slopes * relative_pos.to(dtype=slopes.dtype) @@ -240,6 +244,24 @@ def attention_ref( output: (batch_size, seqlen_q, nheads, head_dim) attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout """ + PRINT_DEBUG = DEBUG and upcast==True + + if PRINT_DEBUG: + print() + print("attention_ref") + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("query_padding_mask:",query_padding_mask) + print("key_padding_mask:",key_padding_mask) + print("attn_bias:",attn_bias) + print("dropout_p:",dropout_p) + print("dropout_mask:",dropout_mask) + print("causal:",causal) + print("window_size:",window_size) + print("upcast:",upcast) + print("reorder_ops:",reorder_ops) + if causal: window_size = (window_size[0], 0) dtype_og = q.dtype @@ -249,12 +271,26 @@ def attention_ref( k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) d = q.shape[-1] - if not reorder_ops: - scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k) + if False: + if not reorder_ops: + scores = torch.einsum("bthd,bshd->bhts", q, k) + else: + scores = torch.einsum("bthd,bshd->bhts", q, k) else: - scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d)) + if not reorder_ops: + scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k) + else: + scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d)) + + + + if key_padding_mask is not None: + if PRINT_DEBUG: + print("scores before key padding mask:", scores, scores.shape) scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")) + if PRINT_DEBUG: + print("scores after key padding mask:", scores, scores.shape) if window_size[0] >= 0 or window_size[1] >= 0: local_mask = construct_local_mask( seqlen_q, @@ -264,9 +300,18 @@ def attention_ref( key_padding_mask, q.device, ) + if PRINT_DEBUG: + print("scores before causal:",scores) scores.masked_fill_(local_mask, float("-inf")) + if PRINT_DEBUG: + print("scores after causal:",scores) if attn_bias is not None: + if PRINT_DEBUG: + print("scores before attn_bias:", scores, scores.shape) scores = scores + attn_bias + if PRINT_DEBUG: + print("scores after attn_bias:", scores, scores.shape) + attention = torch.softmax(scores, dim=-1).to(v.dtype) # Some rows might be completely masked out so we fill them with zero instead of NaN if window_size[0] >= 0 or window_size[1] >= 0: @@ -847,8 +892,6 @@ def is_hip(): def is_power_of_2(n): return n > 0 and (n & (n - 1)) == 0 -DEBUG=False - @pytest.mark.parametrize("kvpacked", [True, False]) # @pytest.mark.parametrize("kvpacked", [False]) @pytest.mark.parametrize("dtype", ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16])) @@ -892,7 +935,7 @@ def test_flash_attn_output( if is_hip(): if dropout_p != 0.0: - pytest.skip("Dropout not supported in HIP") + pytest.skip("Dropout not supported on AMD yet") # skip all cases where seqlen_q, seqlen_k, or d are not powers of 2 if not (is_power_of_2(seqlen_q) and is_power_of_2(seqlen_k) and is_power_of_2(d)): @@ -1184,7 +1227,7 @@ def test_flash_attn_varlen_output( ): if is_hip(): if dropout_p != 0.0: - pytest.skip("Dropout not supported in HIP") + pytest.skip("Dropout not supported on AMD yet") # skip all cases where seqlen_q, seqlen_k, or d are not powers of 2 if not (is_power_of_2(seqlen_q) and is_power_of_2(seqlen_k) and is_power_of_2(d)): @@ -1936,6 +1979,39 @@ def test_flash_attn_kvcache( num_splits, dtype, ): + if DEBUG: + print() + print("test_flash_attn_kvcache") + print("seqlen_q:", seqlen_q) + print("seqlen_k:", seqlen_k ) + print("d:", d ) + print("has_batch_idx:", has_batch_idx ) + print("paged_kv_block_size:", paged_kv_block_size ) + print("rotary_fraction:", rotary_fraction ) + print("rotary_interleaved:", rotary_interleaved ) + print("seqlen_new_eq_seqlen_q:", seqlen_new_eq_seqlen_q ) + print("causal:", causal ) + print("local:", local) + print("alibi:", alibi) + print("new_kv:", new_kv ) + print("mha_type:", mha_type ) + print("num_splits:", num_splits ) + print("dtype:", dtype) + + if is_hip(): + if paged_kv_block_size is not None: + pytest.skip("paged attention not supported on AMD yet") + + if local == True: + pytest.skip("local sliding window attention not supported on AMD yet") + + if rotary_interleaved == True or rotary_fraction > 0.0: + pytest.skip("rotary embedding not supported on AMD yet") + + # skip all cases where seqlen_q, seqlen_k, or d are not powers of 2 + if not (is_power_of_2(seqlen_q) and is_power_of_2(seqlen_k) and is_power_of_2(d)): + pytest.skip("seqlen_q, seqlen_k, or d are not powers of 2") + if seqlen_q > seqlen_k and new_kv: pytest.skip() if not new_kv and rotary_fraction > 0.0: @@ -1948,12 +2024,22 @@ def test_flash_attn_kvcache( batch_size = 2 batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 nheads = 6 + + if DEBUG: + print("nheads_q:", nheads) + print("batch_size:", batch_size) + # rotary_dim must be a multiple of 16, and must be <= d rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) assert nheads % nheads_k == 0 window_size = (-1, -1) if not local else torch.randint(0, seqlen_k, (2,)) - q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) + if False: + q = torch.zeros(batch_size_cache, seqlen_q, nheads_k, d, device=device, dtype=dtype) + for i in range(seqlen_q): + q[:, i, :, :] = torch.full((batch_size_cache, nheads_k, d), i + 1, device=device, dtype=dtype) + else: + q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) seqlen_new = seqlen_q if seqlen_new_eq_seqlen_q else torch.randint(1, seqlen_q + 1, (1,)).item() if new_kv: k = torch.randn(batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype) @@ -1961,9 +2047,37 @@ def test_flash_attn_kvcache( else: k, v = None, None if paged_kv_block_size is None: - k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) - v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + # Increasing Cache + if False: + k_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + + for i in range(nheads_k): + k_cache[:, :, i, :] = torch.full((batch_size_cache, seqlen_k, d), i + 1, device=device, dtype=dtype) + v_cache[:, :, i, :] = torch.full((batch_size_cache, seqlen_k, d), i + 1, device=device, dtype=dtype) + elif False: + k_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + + for i in range(seqlen_k): + k_cache[:, i, :, :] = torch.full((batch_size_cache, nheads_k, d), i + 1, device=device, dtype=dtype) + v_cache[:, i, :, :] = torch.full((batch_size_cache, nheads_k, d), i + 1, device=device, dtype=dtype) + + elif False: + k_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.zeros(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + + values = torch.arange(1, d + 1, device=device, dtype=dtype).view(1, 1, 1, d) + k_cache[:, :, :, :] = values.expand(batch_size_cache, seqlen_k, nheads_k, d) + v_cache[:, :, :, :] = values.expand(batch_size_cache, seqlen_k, nheads_k, d) + else: + k_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) + v_cache = torch.randn(batch_size_cache, seqlen_k, nheads_k, d, device=device, dtype=dtype) block_table = None + if DEBUG: + print("k_cache:", k_cache, k_cache.shape) + print("v_cache:", v_cache, v_cache.shape) + print("block_table:", block_table, block_table.shape if block_table is not None else None) else: ( k_cache, @@ -1975,6 +2089,13 @@ def test_flash_attn_kvcache( ) = _generate_block_kvcache( seqlen_k, paged_kv_block_size, batch_size, nheads_k, d, device, dtype ) + if DEBUG: + print("k_cache_paged:", k_cache_paged, k_cache_paged.shape) + print("v_cache_paged:", v_cache_paged, v_cache_paged.shape) + print("block_table:", block_table, block_table.shape) + print("k_cache:", k_cache, k_cache.shape) + print("v_cache:", v_cache, v_cache.shape) + print("num_blocks:", num_blocks) cache_seqlens = torch.randint( 0 if new_kv else 1, # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough @@ -1994,6 +2115,14 @@ def test_flash_attn_kvcache( ] else: cache_batch_idx = None + if DEBUG: + print("cache_seqlens:", cache_seqlens) + print("arange:", arange) + print("cache_seqlens_expanded:", cache_seqlens_expanded) + print("seqlen_new:", seqlen_new) + print("key_padding_mask:", key_padding_mask) + print("cache_batch_idx:", cache_batch_idx) + if alibi: alibi_slopes = torch.rand(batch_size, nheads, device=device, dtype=torch.float32) * 0.3 attn_bias = attn_bias_from_alibi_slopes( @@ -2105,6 +2234,14 @@ def test_flash_attn_kvcache( upcast=False, reorder_ops=True, ) + + if DEBUG: + print() + print("out:", out, out.shape) + print("out_ref:", out_ref, out_ref.shape) + print("out_pt:", out_pt, out_pt.shape) + print() + print(f"Output max diff: {(out - out_ref).abs().max().item()}") print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") @@ -2131,8 +2268,23 @@ def test_flash_attn_kvcache( "(b nblocks) block_size ... -> b (nblocks block_size) ...", b=batch_size, )[:, :seqlen_k] + + if DEBUG: + print("k_cache_select:", k_cache_select, k_cache_select.shape) + print("k_cache_ref:", k_cache_ref, k_cache_ref.shape) assert torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3) + if DEBUG: + print("v_cache_select:", v_cache_select, v_cache_select.shape) + print("v_cache_ref:", v_cache_ref, v_cache_ref.shape) assert torch.equal(v_cache_select, v_cache_ref) + else: + if DEBUG: + print("k_cache_select:", k_cache, k_cache.shape) + print("k_cache_ref:", k_cache_ref, k_cache_ref.shape) + if DEBUG: + print("v_cache_select:", v_cache, v_cache.shape) + print("v_cache_ref:", v_cache_ref, v_cache_ref.shape) + mult = 3 if not alibi else 5 assert (out - out_ref).abs().max().item() <= mult * (out_pt - out_ref).abs().max().item() + 1e-5