From 5e53c8b7c82f3d8da6f2a1b8620d5950d114f001 Mon Sep 17 00:00:00 2001 From: baoqiwen Date: Mon, 13 Apr 2026 11:36:52 +0800 Subject: [PATCH] Support d != dv test, add varlen test. --- benchmark_fa4_mask_mod.py | 63 +++-- benchmark_fa4_varlen_paddle.py | 449 ++++++++++++++++++++++++++++++++ benchmark_fa4_varlen_torch.py | 457 +++++++++++++++++++++++++++++++++ benchmark_flashmask.py | 55 ++-- plot_radar.py | 12 +- test_flashmask.py | 1 + test_flashmask_varlen.py | 306 ++++++++++++++++++++++ 7 files changed, 1290 insertions(+), 53 deletions(-) create mode 100644 benchmark_fa4_varlen_paddle.py create mode 100644 benchmark_fa4_varlen_torch.py create mode 100644 test_flashmask_varlen.py diff --git a/benchmark_fa4_mask_mod.py b/benchmark_fa4_mask_mod.py index b1e26c8..a0c993a 100644 --- a/benchmark_fa4_mask_mod.py +++ b/benchmark_fa4_mask_mod.py @@ -208,6 +208,7 @@ def test_mask( HKV, S, D, + DV, dtype, skip_correctness: bool = False, print_mask: bool = True, @@ -257,20 +258,21 @@ def test_mask( tile_n=tile_n ) - q, out, gradOut = [ - torch.randn(B, S, H, D, device=device, dtype=data_type, requires_grad=True) - for _ in range(3) - ] - k, v = [ - torch.randn(B, S, HKV, D, device=device, dtype=data_type, requires_grad=True) + q = torch.randn(B, S, H, D, device=device, dtype=data_type, requires_grad=True) + k = torch.randn(B, S, HKV, D, device=device, dtype=data_type, requires_grad=True) + v = torch.randn(B, S, HKV, DV, device=device, dtype=data_type, requires_grad=True) + + out, gradOut = [ + torch.randn(B, S, H, DV, device=device, dtype=data_type, requires_grad=True) for _ in range(2) ] + lse = torch.empty(B, H, S, device=device, dtype=torch.float32) fa4_mask_mod_call = lambda: _flash_attn_fwd( q=q, k=k, v=v, out=out, lse=lse, softmax_scale=None, causal=causal, - m_block_size=tile_m, n_block_size=tile_n, + tile_mn=(tile_m, tile_n), mask_mod=cute_mask_mod, block_sparse_tensors=block_sparse_tensors_fwd, aux_tensors=aux_tensors, return_lse=True, @@ -842,8 +844,13 @@ def main(examples: List[str] = ["all"], dtype='bf16'): #doc_seq_lens_list = doc_seq_lens_list[::-1] # Note(umiswing): fa4 does not support d 256 - for D in [128]: - H = 4096 // D + for D in [128, 192]: + if D == 192: + DV = 128 + H = 16 + else: + DV: int = D + H = 4096 // D HKV = H for idx, (S, prefix_doc_seq_lens, qksparse_mask) in enumerate(doc_seq_lens_list): B = 128 * 1024 // S @@ -851,7 +858,7 @@ def main(examples: List[str] = ["all"], dtype='bf16'): doc_seq_lens = [x[1] for x in prefix_doc_seq_lens] maskout_pair = [] offset = 0 - print(f"{B}_{S}_{H}_{D}_{idx}_{dtype}") + print(f"{B}_{S}_{H}_{D}_{DV}_{idx}_{dtype}") if sum(qksparse_mask) == 0: maskout_pair = [(1024, 538), (2358, 1700)] else: @@ -863,22 +870,28 @@ def main(examples: List[str] = ["all"], dtype='bf16'): share_qa_docs = [split_sequence(doc_seq) for doc_seq in doc_seq_lens] available_examples = { - "Full": lambda: test_mask(mask_info={"cute_mask_mod": None, "flex_mask_mod": None, "aux_tensors": None, "causal": False}, B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal": lambda: test_mask(mask_info={"cute_mask_mod": None, "flex_mask_mod": None, "aux_tensors": None, "causal": True}, B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Sliding Window": lambda: test_mask(mask_info=generate_sliding_window(window_size=int(S*0.0625)), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal Document Mask": lambda: test_mask(mask_info=generate_causal_document_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - # Note(umiswing): FA4 mask_mod will hang in Document Mask, and idk why - # "Document Mask": lambda: test_mask(mask_info=generate_document_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Share Question Mask": lambda: test_mask(mask_info=generate_share_question_mask(doc_seq_lens=share_qa_docs, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Global Sliding Window": lambda: test_mask(mask_info=generate_global_sliding_window_mask(global_token=16, B=B, S=S, window_size=(int(S*0.0625), int(S*0.0625))), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal Blockwise Mask": lambda: test_mask(mask_info=generate_causal_blockwise_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Prefix LM Document Mask": lambda: test_mask(mask_info=generate_prefix_lm_document_mask(doc_seq_lens=prefix_doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Prefix LM Causal Mask": lambda: test_mask(mask_info=generate_prefix_lm_causal_mask(prefix_length=int(S*0.5), B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "QK-sparse Mask": lambda: test_mask(mask_info=generate_qk_sparse_mask(maskout_pair=maskout_pair, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - "Random Eviction Mask": lambda: test_mask(mask_info=generate_random_eviction_mask(start_row=S//2, B=B, S=S, H=H, HKV=HKV), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), - # "Hybrid SWA Prefix LM Doc": lambda: test_mask(mask_info=generate_hybrid_swa_prefix_lm_document_mask(batch_size=B, seqlen=S, hkv=H, d=D, doc_seq_lens=prefix_doc_seq_lens), B=B, S=S, H=H, HKV=HKV, D=D, dtype=dtype), + "Full": lambda: test_mask(mask_info={"cute_mask_mod": None, "flex_mask_mod": None, "aux_tensors": None, "causal": False}, B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal": lambda: test_mask(mask_info={"cute_mask_mod": None, "flex_mask_mod": None, "aux_tensors": None, "causal": True}, B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), } + # d=192, dv=128: mask_mod is not None is not supported, + # only Full and Causal are available. + if D != 192: + available_examples.update({ + "Sliding Window": lambda: test_mask(mask_info=generate_sliding_window(window_size=int(S*0.0625)), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal Document Mask": lambda: test_mask(mask_info=generate_causal_document_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + # Note(umiswing): FA4 mask_mod will hang in Document Mask, and idk why + # "Document Mask": lambda: test_mask(mask_info=generate_document_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Share Question Mask": lambda: test_mask(mask_info=generate_share_question_mask(doc_seq_lens=share_qa_docs, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Global Sliding Window": lambda: test_mask(mask_info=generate_global_sliding_window_mask(global_token=16, B=B, S=S, window_size=(int(S*0.0625), int(S*0.0625))), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal Blockwise Mask": lambda: test_mask(mask_info=generate_causal_blockwise_mask(doc_seq_lens=doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Prefix LM Document Mask": lambda: test_mask(mask_info=generate_prefix_lm_document_mask(doc_seq_lens=prefix_doc_seq_lens, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Prefix LM Causal Mask": lambda: test_mask(mask_info=generate_prefix_lm_causal_mask(prefix_length=int(S*0.5), B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "QK-sparse Mask": lambda: test_mask(mask_info=generate_qk_sparse_mask(maskout_pair=maskout_pair, B=B, S=S), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Random Eviction Mask": lambda: test_mask(mask_info=generate_random_eviction_mask(start_row=S//2, B=B, S=S, H=H, HKV=HKV), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + # "Hybrid SWA Prefix LM Doc": lambda: test_mask(mask_info=generate_hybrid_swa_prefix_lm_document_mask(batch_size=B, seqlen=S, hkv=H, d=D, doc_seq_lens=prefix_doc_seq_lens), B=B, S=S, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + }) + if "all" in examples: ex_to_run = list(available_examples.keys()) else: @@ -916,7 +929,7 @@ def main(examples: List[str] = ["all"], dtype='bf16'): ) content2=tabulate(results, headers=headers, tablefmt="tsv") os.makedirs(f"{dtype}", exist_ok=True) - text_file = open(f"{dtype}/fa4_mask_mod_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{idx}.csv","w") + text_file = open(f"{dtype}/fa4_mask_mod_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}.csv","w") text_file.write(content2) text_file.close() diff --git a/benchmark_fa4_varlen_paddle.py b/benchmark_fa4_varlen_paddle.py new file mode 100644 index 0000000..43d2825 --- /dev/null +++ b/benchmark_fa4_varlen_paddle.py @@ -0,0 +1,449 @@ +""" +Performance benchmark for flash_attn_varlen_func (Paddle varlen). + +Causal Document Mask (FA mask_mod) and flash_attn_varlen are equivalent: + - varlen: packs multiple independent sequences into one flat tensor, + using cu_seqlens to delimit boundaries; causal=True is applied per-sequence. + - Causal Document Mask: a single padded batch where a mask prevents + cross-document attention and enforces causal ordering. + +This script benchmarks paddle flash_attn_varlen_func (FLAGS_flash_attn_version=3) +across a range of configs that mirror the shape / doc-distribution used in +benchmark_flashmask.py, so numbers can be compared directly. + +Usage: + python benchmark_fa4_varlen.py + python benchmark_fa4_varlen.py --dtype fp16 + python benchmark_fa4_varlen.py --seqlens 2048 4096 8192 + python benchmark_fa4_varlen.py --from_file kernel_test_seq_info.txt +""" + +import os +import random +import argparse +from datetime import datetime +from typing import List, Optional, Tuple + +import numpy as np +import paddle +from tabulate import tabulate + +from flash_mask.flash_attn_v4.paddle.interface import _flash_attn_fwd, _flash_attn_bwd + +# --------------------------------------------------------------------------- +# Reproducibility +# --------------------------------------------------------------------------- +paddle.seed(0) +np.random.seed(0) +random.seed(0) +paddle.set_device('gpu') + +# --------------------------------------------------------------------------- +# Benchmarking helpers (mirror benchmark_flashmask.py) +# --------------------------------------------------------------------------- + +def _summarize_statistics(times, quantiles, return_mode): + if quantiles is not None: + ret = paddle.quantile(times, paddle.to_tensor(quantiles, dtype=paddle.float32)).tolist() + if len(ret) == 1: + ret = ret[0] + return ret + if return_mode == "all": + return times.tolist() + return getattr(paddle, return_mode)(times).item() + + +def do_bench(fn, warmup=10, rep=50, grad_to_none=None, + quantiles=None, return_mode="mean"): + assert return_mode in ["min", "max", "mean", "median", "all"] + + fn() + paddle.device.synchronize() + + # L2 cache flush buffer (256 MB) + cache = paddle.empty([int(256 * 1024 * 1024 // 4)], dtype=paddle.int32) + + start_event = [paddle.device.Event(enable_timing=True) for _ in range(rep)] + end_event = [paddle.device.Event(enable_timing=True) for _ in range(rep)] + + for _ in range(warmup): + fn() + + for i in range(rep): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + cache.zero_() + start_event[i].record() + fn() + end_event[i].record() + + paddle.device.synchronize() + times = paddle.to_tensor( + [s.elapsed_time(e) for s, e in zip(start_event, end_event)], + dtype=paddle.float32, + ) + return _summarize_statistics(times, quantiles, return_mode) + + +def cal_flops(B_eff, H, Sq_total, Sk_per_seq, D, mode="fwd"): + """ + Approximate FLOPs for a varlen batch. + + B_eff – effective number of sequences + H – number of query heads + Sq_total – total query tokens + Sk_per_seq – list of per-sequence key lengths (same as query lens for self-attn) + D – head dimension + """ + assert mode in ["fwd", "bwd", "fwd_bwd"] + # Each sequence i contributes 4 * seqlen_i^2 * H * D for causal (density=0.5) + # or 4 * seqlen_i^2 * H * D for full attention. + flops = sum(4 * s * s * H * D for s in Sk_per_seq) + mult = {"fwd": 1.0, "bwd": 2.5, "fwd_bwd": 3.5}[mode] + return mult * flops + + +def cal_tflops(flops, time_ms): + return flops * (1e3 / time_ms) / 1e12 + + +def print_header(text): + width = 91 + print("╔" + "═" * (width - 2) + "╗") + print(f"║ {text.center(width - 4)} ║") + print("╚" + "═" * (width - 2) + "╝") + + +# --------------------------------------------------------------------------- +# varlen tensor builder +# --------------------------------------------------------------------------- + +def build_varlen_tensors( + doc_seq_lens: List[int], + H: int, + HKV: int, + D: int, + DV: int, + dtype, +): + """ + Pack multiple sequences of lengths doc_seq_lens into flat varlen tensors. + + Returns q, k, v (total_tokens, nheads, D/DV) + cu_seqlens_q, cu_seqlens_k (B+1,) int32 + max_seqlen + """ + total = sum(doc_seq_lens) + q = paddle.randn([total, H, D], dtype=dtype) + k = paddle.randn([total, HKV, D], dtype=dtype) + v = paddle.randn([total, HKV, DV], dtype=dtype) + + lens_t = paddle.to_tensor(doc_seq_lens, dtype=paddle.int32) + cu_seqlens = paddle.concat([ + paddle.zeros([1], dtype=paddle.int32), + lens_t.cumsum(0).cast(paddle.int32), + ]) + + max_seqlen = max(doc_seq_lens) + return q, k, v, cu_seqlens, max_seqlen + + +# --------------------------------------------------------------------------- +# Causal Document Mask → varlen equivalence conversion +# --------------------------------------------------------------------------- + +def causal_doc_mask_to_varlen( + doc_seq_lens: List[int], + S: int, +) -> Tuple[List[int], int]: + """ + A Causal Document Mask over a padded sequence of length S with + sub-sequence lengths doc_seq_lens is *equivalent* to running + flash_attn_varlen with causal=True over those same doc_seq_lens. + + This function just returns the canonical (doc_seq_lens, max_seqlen) + tuple that varlen needs, stripping any padding from the last doc. + """ + # doc_seq_lens already sums to S (including padding absorbed into last doc) + assert sum(doc_seq_lens) == S, ( + f"sum(doc_seq_lens)={sum(doc_seq_lens)} != S={S}" + ) + max_seqlen = max(doc_seq_lens) + return doc_seq_lens, max_seqlen + + +# --------------------------------------------------------------------------- +# Core benchmark function +# --------------------------------------------------------------------------- + +def benchmark_varlen( + doc_seq_lens: List[int], + S: int, + H: int, + HKV: int, + D: int, + DV: int, + dtype, + causal: bool = True, +): + """ + Benchmark flash_attn_varlen_func fwd + bwd for the given doc distribution. + Equivalent to the Causal Document Mask benchmark in benchmark_flashmask.py. + """ + seq_lens, max_seqlen = causal_doc_mask_to_varlen(doc_seq_lens, S) + + q, k, v, cu_seqlens, max_seqlen = build_varlen_tensors( + seq_lens, H, HKV, D, DV, dtype + ) + + softmax_scale = D ** -0.5 + + # ---------- forward ---------- + # Call _flash_attn_fwd directly — no autograd graph, pure kernel timing. + fwd_fn = lambda: _flash_attn_fwd( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + return_lse=True, + causal=causal, + softmax_scale=softmax_scale, + ) + + fwd_time_ms = do_bench(fwd_fn) + + # ---------- backward ---------- + # Run fwd once to get out and lse, then benchmark bwd kernel directly. + # _flash_attn_bwd is a pure function: same inputs → same gradients, + # no autograd state, no graph, no PyLayer overhead. + out, lse = fwd_fn() + grad_out = paddle.randn_like(out) + + bwd_fn = lambda: _flash_attn_bwd( + q, k, v, + out=out, + dout=grad_out, + lse=lse, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + causal=causal, + softmax_scale=softmax_scale, + ) + + bwd_time_ms = do_bench(bwd_fn) + + total_time_ms = fwd_time_ms + bwd_time_ms + + # ---------- FLOPs (causal ⇒ density ≈ 0.5 per-seq) ---------- + # FLOPs formula: 4 * seqlen^2 * H * D (D here is head_dim of QK, i.e. D) + # matches benchmark_flashmask.py cal_flops convention + density = 0.5 if causal else 1.0 + fwd_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "fwd") + bwd_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "bwd") + total_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "fwd_bwd") + + fwd_tflops = cal_tflops(fwd_flops, fwd_time_ms) + bwd_tflops = cal_tflops(bwd_flops, bwd_time_ms) + total_tflops = cal_tflops(total_flops, total_time_ms) + sparsity = 1.0 - density + + return (fwd_time_ms, bwd_time_ms, total_time_ms, + fwd_flops, bwd_flops, total_flops, + fwd_tflops, bwd_tflops, total_tflops, + sparsity) + + +# --------------------------------------------------------------------------- +# Doc-len generators (mirrors benchmark_flashmask.py helpers) +# --------------------------------------------------------------------------- + +def make_uniform_docs(S: int, n_docs: int) -> List[int]: + """Divide S tokens into n_docs equal-ish chunks.""" + base = S // n_docs + lens = [base] * n_docs + lens[-1] += S - sum(lens) + return lens + + +def make_random_docs(S: int, n_docs: int, seed: int = 42) -> List[int]: + """Random split of S into n_docs positive chunks.""" + rng = np.random.default_rng(seed) + cuts = np.sort(rng.choice(S - 1, n_docs - 1, replace=False)) + 1 + cuts = np.concatenate([[0], cuts, [S]]) + lens = [int(cuts[i+1] - cuts[i]) for i in range(n_docs)] + return lens + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main( + seqlens: List[int] = None, + D: int = 192, + DV: int = 128, + dtype: str = "bf16", + from_file: Optional[str] = None, + n_docs: int = 4, + causal: bool = True, +): + if seqlens is None: + seqlens = [2048, 4096, 8192] + + paddle.set_flags({'FLAGS_flash_attn_version': 3}) + + dtype_map = {"bf16": paddle.bfloat16, "fp16": paddle.float16} + data_type = dtype_map[dtype] + + current_time = datetime.now().strftime("%Y%m%d_%H%M%S") + os.makedirs(dtype, exist_ok=True) + + # ---------------------------------------------------------------- + # Build list of (S, doc_seq_lens) configs + # ---------------------------------------------------------------- + configs: List[Tuple[int, List[int]]] = [] + + if from_file is not None: + # Same format as benchmark_flashmask.py reads kernel_test_seq_info.txt + total_length = 0 + with open(from_file, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + if "Total length" in line: + total_length = int(line.split(":")[1].split(",")[0].strip()) + else: + doc_list = eval(line.split(":")[-1].split("#")[0].strip()) + doc_seq_lens = [x[1] for x in doc_list] + configs.append((total_length, doc_seq_lens)) + else: + for S in seqlens: + configs.append((S, make_random_docs(S, n_docs))) + + # ---------------------------------------------------------------- + # Run benchmarks + # ---------------------------------------------------------------- + for D in [128, 192, 256]: + if D == 192: + DV = 128 + H = 16 + else: + DV: int = D + H = 4096 // D + + HKV = H + + for idx, (S, doc_seq_lens) in enumerate(configs): + B = 128 * 1024 // S # match benchmark_flashmask convention + + print(f"{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}_{dtype}") + # print(f" doc_seq_lens = {doc_seq_lens}") + + # Replicate doc_seq_lens B times so total tokens = B*S, + # matching benchmark_flashmask which processes B padded samples. + batched_doc_seq_lens = doc_seq_lens * B + + results = [] + + # ---- varlen ---- + label = "varlen causal" + print(label) + try: + (fwd_t, bwd_t, tot_t, + fwd_f, bwd_f, tot_f, + fwd_tf, bwd_tf, tot_tf, + sparsity) = benchmark_varlen( + doc_seq_lens=batched_doc_seq_lens, + S=S * B, H=H, HKV=HKV, D=D, DV=DV, + dtype=data_type, causal=causal, + ) + results.append([ + label, + f"{fwd_t:.4f}", + f"{bwd_t:.4f}", + f"{tot_t:.4f}", + f"{fwd_f:.4f}", + f"{bwd_f:.4f}", + f"{tot_f:.4f}", + f"{fwd_tf:.4f}", + f"{bwd_tf:.4f}", + f"{tot_tf:.4f}", + f"{sparsity:.4f}", + ]) + except Exception as e: + print(f" ERROR: {e}") + results.append([label] + ["ERROR"] * 10) + + headers = [ + "Operation", + "FW Time (ms)", + "BW Time (ms)", + "TOTAL Time (ms)", + "FW FLOPs", + "BW FLOPs", + "TOTAL FLOPs", + "FW TFLOPs/s", + "BW TFLOPs/s", + "TOTAL TFLOPs/s", + "Sparsity", + ] + print( + tabulate( + results, + headers=headers, + tablefmt="grid", + ) + ) + content2 = tabulate(results, headers=headers, tablefmt="tsv") + os.makedirs(dtype, exist_ok=True) + text_file = open(f"{dtype}/fa4_varlen_paddle_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}.csv", "w") + text_file.write(content2) + text_file.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Benchmark paddle flash_attn_varlen_func (Paddle varlen).\n" + "Causal Document Mask and varlen are equivalent — see module docstring." + ) + parser.add_argument( + "--seqlens", type=int, nargs="+", default=[2048, 4096, 8192], + help="List of total sequence lengths to benchmark.", + ) + parser.add_argument( + "--D", type=int, default=128, + help="Query/Key head dimension.", + ) + parser.add_argument( + "--DV", type=int, default=128, + help="Value head dimension.", + ) + parser.add_argument( + "--dtype", type=str, default="bf16", choices=["bf16", "fp16"], + help="Data type.", + ) + parser.add_argument( + "--from_file", type=str, default="kernel_test_seq_info.txt", + help="Read doc-seq-len configs from a kernel_test_seq_info.txt file " + "(same format as benchmark_flashmask.py).", + ) + parser.add_argument( + "--n_docs", type=int, default=4, + help="Number of documents to split each sequence into (random split).", + ) + parser.add_argument( + "--no_causal", action="store_true", + help="Disable causal masking (runs full attention instead).", + ) + args = parser.parse_args() + + main( + seqlens=args.seqlens, + D=args.D, + DV=args.DV, + dtype=args.dtype, + from_file=args.from_file, + n_docs=args.n_docs, + causal=not args.no_causal, + ) diff --git a/benchmark_fa4_varlen_torch.py b/benchmark_fa4_varlen_torch.py new file mode 100644 index 0000000..691aea6 --- /dev/null +++ b/benchmark_fa4_varlen_torch.py @@ -0,0 +1,457 @@ +""" +Performance benchmark for flash_attn_varlen_func (FA4 varlen). + +Causal Document Mask (FA4 mask_mod) and flash_attn_varlen are equivalent: + - varlen: packs multiple independent sequences into one flat tensor, + using cu_seqlens to delimit boundaries; causal=True is applied per-sequence. + - Causal Document Mask: a single padded batch where a mask prevents + cross-document attention and enforces causal ordering. + +This script benchmarks flash_attn_varlen_func across a range of configs +that mirror the shape / doc-distribution used in benchmark_fa4_mask_mod.py, +so numbers can be compared directly. + +Usage: + python tests/cute/benchmark_varlen.py + python tests/cute/benchmark_varlen.py --dtype fp16 + python tests/cute/benchmark_varlen.py --seqlens 2048 4096 8192 --D 128 + python tests/cute/benchmark_varlen.py --from_file kernel_test_seq_info.txt +""" + +import os +import sys +import random +import argparse +from datetime import datetime +from typing import List, Optional, Tuple + +import numpy as np +import torch +from tabulate import tabulate + +from flash_attn.cute.interface import _flash_attn_fwd, _flash_attn_bwd +# from flash_mask.flash_attn_v4.torch.interface import _flash_attn_fwd, _flash_attn_bwd + +# --------------------------------------------------------------------------- +# Reproducibility +# --------------------------------------------------------------------------- +torch.manual_seed(0) +np.random.seed(0) +random.seed(0) +torch.set_default_device("cuda") + +# --------------------------------------------------------------------------- +# Benchmarking helpers (mirror benchmark_fa4_mask_mod.py) +# --------------------------------------------------------------------------- + +def _summarize_statistics(times, quantiles, return_mode): + if quantiles is not None: + ret = torch.quantile(times, torch.tensor(quantiles, dtype=torch.float32)).tolist() + if len(ret) == 1: + ret = ret[0] + return ret + if return_mode == "all": + return times.tolist() + return getattr(torch, return_mode)(times).item() + + +def do_bench(fn, warmup=10, rep=50, grad_to_none=None, + quantiles=None, return_mode="mean"): + assert return_mode in ["min", "max", "mean", "median", "all"] + + fn() + torch.cuda.synchronize() + + # L2 cache flush buffer (256 MB) + cache = torch.empty(int(256 * 1024 * 1024 // 4), dtype=torch.int32, device="cuda") + + start_event = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + end_event = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + + for _ in range(warmup): + fn() + + for i in range(rep): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + cache.zero_() + start_event[i].record() + fn() + end_event[i].record() + + torch.cuda.synchronize() + times = torch.tensor( + [s.elapsed_time(e) for s, e in zip(start_event, end_event)], + dtype=torch.float32, + ) + return _summarize_statistics(times, quantiles, return_mode) + + +def cal_flops(B_eff, H, Sq_total, Sk_per_seq, D, mode="fwd"): + """ + Approximate FLOPs for a varlen batch. + + B_eff – effective number of sequences + H – number of query heads + Sq_total – total query tokens + Sk_per_seq – list of per-sequence key lengths (same as query lens for self-attn) + D – head dimension + """ + assert mode in ["fwd", "bwd", "fwd_bwd"] + # Each sequence i contributes 4 * seqlen_i^2 * H * D for causal (density=0.5) + # or 4 * seqlen_i^2 * H * D for full attention. + flops = sum(4 * s * s * H * D for s in Sk_per_seq) + mult = {"fwd": 1.0, "bwd": 2.5, "fwd_bwd": 3.5}[mode] + return mult * flops + + +def cal_tflops(flops, time_ms): + return flops * (1e3 / time_ms) / 1e12 + + +def print_header(text): + width = 91 + print("╔" + "═" * (width - 2) + "╗") + print(f"║ {text.center(width - 4)} ║") + print("╚" + "═" * (width - 2) + "╝") + + +# --------------------------------------------------------------------------- +# varlen tensor builder +# --------------------------------------------------------------------------- + +def build_varlen_tensors( + doc_seq_lens: List[int], + H: int, + HKV: int, + D: int, + DV: int, + dtype: torch.dtype, + device: str = "cuda", +): + """ + Pack multiple sequences of lengths doc_seq_lens into flat varlen tensors. + + Returns q, k, v (total_tokens, nheads, D/DV) + cu_seqlens_q, cu_seqlens_k (B+1,) int32 + max_seqlen + """ + total = sum(doc_seq_lens) + q = torch.randn(total, H, D, device=device, dtype=dtype, requires_grad=True) + k = torch.randn(total, HKV, D, device=device, dtype=dtype, requires_grad=True) + v = torch.randn(total, HKV, DV, device=device, dtype=dtype, requires_grad=True) + + lens_t = torch.tensor(doc_seq_lens, dtype=torch.int32, device=device) + cu_seqlens = torch.cat( + [torch.zeros(1, dtype=torch.int32, device=device), + lens_t.cumsum(0).to(torch.int32)] + ) + + max_seqlen = max(doc_seq_lens) + return q, k, v, cu_seqlens, max_seqlen + + +# --------------------------------------------------------------------------- +# Causal Document Mask → varlen equivalence conversion +# --------------------------------------------------------------------------- + +def causal_doc_mask_to_varlen( + doc_seq_lens: List[int], + S: int, +) -> Tuple[List[int], int]: + """ + A Causal Document Mask over a padded sequence of length S with + sub-sequence lengths doc_seq_lens is *equivalent* to running + flash_attn_varlen with causal=True over those same doc_seq_lens. + + This function just returns the canonical (doc_seq_lens, max_seqlen) + tuple that varlen needs, stripping any padding from the last doc. + """ + # doc_seq_lens already sums to S (including padding absorbed into last doc) + assert sum(doc_seq_lens) == S, ( + f"sum(doc_seq_lens)={sum(doc_seq_lens)} != S={S}" + ) + max_seqlen = max(doc_seq_lens) + return doc_seq_lens, max_seqlen + + +# --------------------------------------------------------------------------- +# Core benchmark function +# --------------------------------------------------------------------------- + +def benchmark_varlen( + doc_seq_lens: List[int], + S: int, + H: int, + HKV: int, + D: int, + DV: int, + dtype: torch.dtype, + causal: bool = True, + device: str = "cuda", +): + """ + Benchmark flash_attn_varlen_func fwd + bwd for the given doc distribution. + Equivalent to the Causal Document Mask benchmark in benchmark_fa4_mask_mod.py. + """ + seq_lens, max_seqlen = causal_doc_mask_to_varlen(doc_seq_lens, S) + + q, k, v, cu_seqlens, max_seqlen = build_varlen_tensors( + seq_lens, H, HKV, D, DV, dtype, device + ) + + softmax_scale = D ** -0.5 + + # ---------- forward ---------- + # Call _flash_attn_fwd directly — no autograd graph, pure kernel timing. + fwd_fn = lambda: _flash_attn_fwd( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + return_lse=True, + causal=causal, + softmax_scale=softmax_scale, + ) + + fwd_time_ms = do_bench(fwd_fn) + + # ---------- backward ---------- + # Run fwd once to get out and lse, then benchmark bwd kernel directly. + # _flash_attn_bwd is a pure function: same inputs → same gradients, + # no autograd state, no graph overhead. + out, lse = fwd_fn() + grad_out = torch.randn_like(out) + + bwd_fn = lambda: _flash_attn_bwd( + q, k, v, + out=out, + dout=grad_out, + lse=lse, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + causal=causal, + softmax_scale=softmax_scale, + ) + + bwd_time_ms = do_bench(bwd_fn) + + total_time_ms = fwd_time_ms + bwd_time_ms + + # ---------- FLOPs (causal ⇒ density ≈ 0.5 per-seq) ---------- + # FLOPs formula: 4 * seqlen^2 * H * D (D here is head_dim of QK, i.e. D) + # matches benchmark_fa4_mask_mod.py cal_flops convention + density = 0.5 if causal else 1.0 + fwd_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "fwd") + bwd_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "bwd") + total_flops = density * cal_flops(len(seq_lens), H, sum(seq_lens), seq_lens, D, "fwd_bwd") + + fwd_tflops = cal_tflops(fwd_flops, fwd_time_ms) + bwd_tflops = cal_tflops(bwd_flops, bwd_time_ms) + total_tflops = cal_tflops(total_flops, total_time_ms) + sparsity = 1.0 - density + + return (fwd_time_ms, bwd_time_ms, total_time_ms, + fwd_flops, bwd_flops, total_flops, + fwd_tflops, bwd_tflops, total_tflops, + sparsity) + + +# --------------------------------------------------------------------------- +# Doc-len generators (mirrors benchmark_fa4_mask_mod.py helpers) +# --------------------------------------------------------------------------- + +def make_uniform_docs(S: int, n_docs: int) -> List[int]: + """Divide S tokens into n_docs equal-ish chunks.""" + base = S // n_docs + lens = [base] * n_docs + lens[-1] += S - sum(lens) + return lens + + +def make_random_docs(S: int, n_docs: int, seed: int = 42) -> List[int]: + """Random split of S into n_docs positive chunks.""" + rng = np.random.default_rng(seed) + cuts = np.sort(rng.choice(S - 1, n_docs - 1, replace=False)) + 1 + cuts = np.concatenate([[0], cuts, [S]]) + lens = [int(cuts[i+1] - cuts[i]) for i in range(n_docs)] + return lens + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main( + seqlens: List[int] = None, + D: int = 192, + DV: int = 128, + dtype: str = "bf16", + from_file: Optional[str] = None, + n_docs: int = 4, + causal: bool = True, +): + if seqlens is None: + seqlens = [2048, 4096, 8192] + + dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16} + data_type = dtype_map[dtype] + + current_time = datetime.now().strftime("%Y%m%d_%H%M%S") + os.makedirs(dtype, exist_ok=True) + + # ---------------------------------------------------------------- + # Build list of (S, doc_seq_lens) configs + # ---------------------------------------------------------------- + configs: List[Tuple[int, List[int]]] = [] + + if from_file is not None: + # Same format as benchmark_fa4_mask_mod.py reads kernel_test_seq_info.txt + total_length = 0 + with open(from_file, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + if "Total length" in line: + total_length = int(line.split(":")[1].split(",")[0].strip()) + else: + doc_list = eval(line.split(":")[-1].split("#")[0].strip()) + doc_seq_lens = [x[1] for x in doc_list] + configs.append((total_length, doc_seq_lens)) + else: + for S in seqlens: + configs.append((S, make_random_docs(S, n_docs))) + + # ---------------------------------------------------------------- + # Run benchmarks + # ---------------------------------------------------------------- + + # for D in [128, 192, 256]: + for D in [128, 192]: + if D == 192: + DV = 128 + H = 16 + else: + DV: int = D + H = 4096 // D + + HKV = H + + for idx, (S, doc_seq_lens) in enumerate(configs): + B = 128 * 1024 // S # match benchmark_fa4_mask_mod convention + + print(f"{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}_{dtype}") + # print(f" doc_seq_lens = {doc_seq_lens}") + + # Replicate doc_seq_lens B times so total tokens = B*S, + # matching benchmark_fa4_mask_mod which processes B padded samples. + batched_doc_seq_lens = doc_seq_lens * B + + results = [] + + # ---- FA4 varlen ---- + label = "varlen causal" + print(label) + try: + (fwd_t, bwd_t, tot_t, + fwd_f, bwd_f, tot_f, + fwd_tf, bwd_tf, tot_tf, + sparsity) = benchmark_varlen( + doc_seq_lens=batched_doc_seq_lens, + S=S * B, H=H, HKV=HKV, D=D, DV=DV, + dtype=data_type, causal=causal, + ) + results.append([ + label, + f"{fwd_t:.4f}", + f"{bwd_t:.4f}", + f"{tot_t:.4f}", + f"{fwd_f:.4f}", + f"{bwd_f:.4f}", + f"{tot_f:.4f}", + f"{fwd_tf:.4f}", + f"{bwd_tf:.4f}", + f"{tot_tf:.4f}", + f"{sparsity:.4f}", + ]) + except Exception as e: + print(f" ERROR: {e}") + results.append([label] + ["ERROR"] * 10) + + headers = [ + "Operation", + "FW Time (ms)", + "BW Time (ms)", + "TOTAL Time (ms)", + "FW FLOPs", + "BW FLOPs", + "TOTAL FLOPs", + "FW TFLOPs/s", + "BW TFLOPs/s", + "TOTAL TFLOPs/s", + "Sparsity", + ] + print( + tabulate( + results, + headers=headers, + tablefmt="grid", + ) + ) + content2 = tabulate(results, headers=headers, tablefmt="tsv") + os.makedirs(dtype, exist_ok=True) + text_file = open(f"{dtype}/fa4_varlen_torch_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}.csv", "w") + text_file.write(content2) + text_file.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Benchmark flash_attn_varlen_func (FA4 varlen).\n" + "Causal Document Mask and varlen are equivalent — see module docstring." + ) + parser.add_argument( + "--seqlens", type=int, nargs="+", default=[2048, 4096, 8192], + help="List of total sequence lengths to benchmark.", + ) + parser.add_argument( + "--D", type=int, default=128, + help="Query/Key head dimension.", + ) + parser.add_argument( + "--DV", type=int, default=128, + help="Value head dimension.", + ) + parser.add_argument( + "--dtype", type=str, default="bf16", choices=["bf16", "fp16"], + help="Data type.", + ) + parser.add_argument( + "--from_file", type=str, default="kernel_test_seq_info.txt", + help="Read doc-seq-len configs from a kernel_test_seq_info.txt file " + "(same format as benchmark_fa4_mask_mod.py).", + ) + parser.add_argument( + "--n_docs", type=int, default=4, + help="Number of documents to split each sequence into (random split).", + ) + parser.add_argument( + "--no_causal", action="store_true", + help="Disable causal masking (runs full attention instead).", + ) + args = parser.parse_args() + + main( + seqlens=args.seqlens, + D=args.D, + DV=args.DV, + dtype=args.dtype, + from_file=args.from_file, + n_docs=args.n_docs, + causal=not args.no_causal, + ) diff --git a/benchmark_flashmask.py b/benchmark_flashmask.py index 685483a..a4146af 100644 --- a/benchmark_flashmask.py +++ b/benchmark_flashmask.py @@ -116,6 +116,7 @@ def test_mask( H, HKV, D, + DV, dtype = 'bf16', ): @@ -126,8 +127,8 @@ def test_mask( query = paddle.randn([B, S, H, D], dtype=data_type) key = paddle.randn([B, SKV, HKV, D], dtype=data_type) - value = paddle.randn([B, SKV, HKV, D], dtype=data_type) - gradOut = paddle.randn([B, S, H, D], dtype=data_type) + value = paddle.randn([B, SKV, HKV, DV], dtype=data_type) + gradOut = paddle.randn([B, S, H, DV], dtype=data_type) query.stop_gradient = False key.stop_gradient = False @@ -754,9 +755,15 @@ def main(examples: List[str] = ["all"], dtype='bf16', fm_version=1, suffix="_bas qksparse_mask = eval(line.split(":")[-1].split("#")[1].strip()) doc_seq_lens_list.append((total_length, doc_list, qksparse_mask)) #doc_seq_lens_list = doc_seq_lens_list[::-1] - for D in [128] if fm_version == 4 else [64, 128, 256]: - H = 4096 // D + for D in [128, 192] if fm_version == 4 else [64, 128, 256]: + if D == 192: + DV = 128 + H = 16 + else: + DV = D + H = 4096 // D HKV = H + for idx, (S, prefix_doc_seq_lens, qksparse_mask) in enumerate(doc_seq_lens_list): B = 128 * 1024 // S @@ -765,10 +772,10 @@ def main(examples: List[str] = ["all"], dtype='bf16', fm_version=1, suffix="_bas doc_seq_lens = [x[1] for x in prefix_doc_seq_lens] maskout_pair = [] offset = 0 - print(f"{B}_{S}_{H}_{HKV}_{D}_{idx}_{dtype}") + print(f"{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}_{dtype}") if not overwrite: - if os.path.exists(f"{dtype}{suffix}/flashmaskv{fm_version}_{B}_{S}_{H}_{D}_{idx}.csv"): - print(f"{dtype}{suffix}/flashmaskv{fm_version}_{B}_{S}_{H}_{D}_{idx}.csv already exists, skipping. To enable overwrite, use: --overwrite (True by default).") + if os.path.exists(f"{dtype}{suffix}/flashmaskv{fm_version}_{B}_{S}_{H}_{D}_{DV}_{idx}.csv"): + print(f"{dtype}{suffix}/flashmaskv{fm_version}_{B}_{S}_{H}_{D}_{DV}_{idx}.csv already exists, skipping. To enable overwrite, use: --overwrite (True by default).") continue if sum(qksparse_mask) == 0: maskout_pair = [(1024, 538), (2358, 1700)] @@ -780,23 +787,22 @@ def main(examples: List[str] = ["all"], dtype='bf16', fm_version=1, suffix="_bas share_qa_docs = [split_sequence(doc_seq) for doc_seq in doc_seq_lens] available_examples = { - "Full": lambda: test_mask(generate_mask_fn=partial(generate_none_mask, causal=False), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal": lambda: test_mask(generate_mask_fn=partial(generate_none_mask, causal=True), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Sliding Window": lambda: test_mask(generate_mask_fn=partial(generate_sliding_window_mask, window_size=int(S*0.0625)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_causal_document_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_document_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Share Question Mask": lambda: test_mask(generate_mask_fn=partial(generate_share_question_mask, doc_seq_lens=share_qa_docs), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - # "Global Sliding Window": lambda: test_mask(generate_mask_fn=partial(generate_global_sliding_window_mask, global_token=16, window_size=(int(S*0.0625), int(S*0.0625))), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Causal Blockwise Mask": lambda: test_mask(generate_mask_fn=partial(generate_causal_blockwise_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Prefix LM Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_prefix_lm_document_mask, doc_seq_lens=prefix_doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Prefix LM Causal Mask": lambda: test_mask(generate_mask_fn=partial(generate_prefix_lm_causal_mask, prefix_length=int(S*0.5)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "QK-sparse Mask": lambda: test_mask(generate_mask_fn=partial(generate_qk_sparse_mask, maskout_pair=maskout_pair), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - "Random Eviction Mask": lambda: test_mask(generate_mask_fn=partial(generate_random_eviction_mask, start_row=S//2), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - # "Hybrid SWA Prefix LM Doc": lambda: test_mask(generate_mask_fn=partial(generate_hybrid_swa_prefix_lm_document_mask, doc_seq_lens=prefix_doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - + "Full": lambda: test_mask(generate_mask_fn=partial(generate_none_mask, causal=False), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal": lambda: test_mask(generate_mask_fn=partial(generate_none_mask, causal=True), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Sliding Window": lambda: test_mask(generate_mask_fn=partial(generate_sliding_window_mask, window_size=int(S*0.0625)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_causal_document_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_document_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Share Question Mask": lambda: test_mask(generate_mask_fn=partial(generate_share_question_mask, doc_seq_lens=share_qa_docs), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + # "Global Sliding Window": lambda: test_mask(generate_mask_fn=partial(generate_global_sliding_window_mask, global_token=16, window_size=(int(S*0.0625), int(S*0.0625))), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Causal Blockwise Mask": lambda: test_mask(generate_mask_fn=partial(generate_causal_blockwise_mask, doc_seq_lens=doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Prefix LM Document Mask": lambda: test_mask(generate_mask_fn=partial(generate_prefix_lm_document_mask, doc_seq_lens=prefix_doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Prefix LM Causal Mask": lambda: test_mask(generate_mask_fn=partial(generate_prefix_lm_causal_mask, prefix_length=int(S*0.5)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "QK-sparse Mask": lambda: test_mask(generate_mask_fn=partial(generate_qk_sparse_mask, maskout_pair=maskout_pair), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + "Random Eviction Mask": lambda: test_mask(generate_mask_fn=partial(generate_random_eviction_mask, start_row=S//2), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + # "Hybrid SWA Prefix LM Doc": lambda: test_mask(generate_mask_fn=partial(generate_hybrid_swa_prefix_lm_document_mask, doc_seq_lens=prefix_doc_seq_lens), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), # Note(umiswing): support load mask and hybrid mask like this, and also, support simulate cp benchmark - # "Dumped Mask": lambda: test_mask(generate_mask_fn=partial(load_mask, path=mask_path, causal=False, cp_size=cp_size, cp_rank=cp_rank), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), - # "Hybrid SWA": lambda: test_mask(generate_mask_fn=partial(load_mask, path=mask_path, causal=False, cp_size=cp_size, cp_rank=cp_rank, hybrid_mask_fn=partial(hybrid_swa, window_size=512, swa_ratio=0.75)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, dtype=dtype), + # "Dumped Mask": lambda: test_mask(generate_mask_fn=partial(load_mask, path=mask_path, causal=False, cp_size=cp_size, cp_rank=cp_rank), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), + # "Hybrid SWA": lambda: test_mask(generate_mask_fn=partial(load_mask, path=mask_path, causal=False, cp_size=cp_size, cp_rank=cp_rank, hybrid_mask_fn=partial(hybrid_swa, window_size=512, swa_ratio=0.75)), B=B, S=SQ, SKV=SKV, H=H, HKV=HKV, D=D, DV=DV, dtype=dtype), } if "all" in examples: @@ -837,8 +843,7 @@ def main(examples: List[str] = ["all"], dtype='bf16', fm_version=1, suffix="_bas content2=tabulate(results, headers=headers, tablefmt="tsv") os.makedirs(f"{dtype}{suffix}", exist_ok=True) # Note(umiswing): this file name is better, but i need to keep the old name for fig plotting - # text_file = open(f"{dtype}{suffix}/flashmaskv{fm_version}_{current_time}_{B}_{SQ}_{SKV}_{H}_{HKV}_{D}_{idx}.csv","w") - text_file = open(f"{dtype}{suffix}/flashmaskv{fm_version}_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{idx}.csv","w") + text_file = open(f"{dtype}{suffix}/flashmaskv{fm_version}_{current_time}_{B}_{S}_{H}_{HKV}_{D}_{DV}_{idx}.csv","w") text_file.write(content2) text_file.close() diff --git a/plot_radar.py b/plot_radar.py index 0be9a38..ca630c0 100644 --- a/plot_radar.py +++ b/plot_radar.py @@ -162,15 +162,22 @@ def main(methods: list = ["flashmaskv1", "flashmaskv3"]): # for kernel in ["fwd", "bwd", "total", "fwd_time", "bwd_time", "total_time", "sparsity"]: for kernel in ["fwd", "bwd", "total"]: for dtype in ['bf16']: - for headdim in [64, 128, 256]: + for headdim in [64, 128, 192, 256]: + headdim_v = 128 if headdim == 192 else headdim categories = {} for seqlen in [8192, 32768, 131072]: method_to_df = {} for method in methods: - filenames = glob.glob(f'{root_dir}/{dtype}/{method}_*{seqlen}_*_{headdim}*.csv') + filenames = glob.glob(f'{root_dir}/{dtype}/{method}_*{seqlen}_*_{headdim}_{headdim_v}*.csv') dataframes = [] for file_path in filenames: df = read_tsv_to_dataframe(file_path) + df.columns = df.columns.str.strip() + # d=192, dv=128: fa4_mask_mod only supports Full & Causal + # (mask_mod is not None is not supported), so truncate all + # methods to the first 2 rows to keep labels aligned. + if headdim == 192 and any(m.startswith('fa4_mask_mod') for m in methods): + df = df[:2] if df is not None: dataframes.append(df) @@ -257,7 +264,6 @@ def main(methods: list = ["flashmaskv1", "flashmaskv3"]): if __name__ == "__main__": from jsonargparse import ArgumentParser parser = ArgumentParser(description="Run specific examples or all examples.") - parser.add_argument( "--methods", type=str, diff --git a/test_flashmask.py b/test_flashmask.py index 362c2df..3d969e3 100644 --- a/test_flashmask.py +++ b/test_flashmask.py @@ -88,6 +88,7 @@ def generate_shapes(): (64, 64), (80, 80), (128, 128), + (192, 128), (192, 192), (256, 256), ]) diff --git a/test_flashmask_varlen.py b/test_flashmask_varlen.py new file mode 100644 index 0000000..a2969b8 --- /dev/null +++ b/test_flashmask_varlen.py @@ -0,0 +1,306 @@ +""" +Cross-framework Flash Attention varlen test. + +Reference : PyTorch F.scaled_dot_product_attention (torch_flash_ref) +Under test: Paddle flash_attn_varlen_func (flash_mask) + +Data is generated via NumPy so that both frameworks receive bit-identical +inputs and upstream gradients. +""" + +# 为什么 ref 使用 PyTorch 而非 Paddle： +# 1. Paddle 的 scaled_dot_product_attention(is_causal=True) 走 memory_efficient_attention (CUTLASS) 路径， +# 不支持 head_dim=256 + GQA/MQA（Q/K 头数不同）的组合，会报 kernel_launched=false。 +# 2. 手写的 matmul+softmax ref 无法精确复现 Flash Attention kernel 的混合精度反向计算 +# （online softmax、分块累加、P→bf16 截断等），在 head_dim=256 时约有 1% 的 case 精度不达标。 +# 3. PyTorch 的 F.scaled_dot_product_attention 支持此类场景, ref 结果与 FA4 varlen 一致。 +# 待 Paddle SDPA 支持上述场景后可切换回来。 + +import math +from typing import Optional + +import numpy as np +import pytest + +import torch +import torch.nn.functional as F_torch + +import paddle +from flash_mask import flash_attn_varlen_func + + +# --------------------------------------------------------------------------- +# Test configuration (edit here to control pytest parametrisation) +# --------------------------------------------------------------------------- + +TEST_B = [1, 7, 20] +TEST_H = [1, 4, 6] +TEST_D = [64, 128, 192, 256] +TEST_MIN_SEQ_LEN = [1, 32, 128] +TEST_MAX_SEQ_LEN = [8, 64, 2048] +TEST_CAUSAL = [True, False] +TEST_SOFTMAX_SCALE = [None, 0.1] +TEST_DTYPE = ["bfloat16", "float16"] +TEST_MHA_TYPE = ["mha", "mqa", "gqa"] + + +# --------------------------------------------------------------------------- +# NumPy data generation (framework-agnostic, bit-reproducible) +# --------------------------------------------------------------------------- + +def generate_varlen_data( + batch_size: int = 8, + n_heads: int = 16, + d_head: int = 128, + min_len: int = 32, + max_len: int = 64, + mha_type: str = "mha", + dtype_str: str = "bfloat16", + seed: int = 0, +): + """Return NumPy arrays for Q, K, V, cu_seqlens, and a random grad_out. + + All floating-point arrays are stored as float32 (bf16 values are first + generated in float32, then rounded to bf16 via a torch round-trip so the + bit patterns are exact bf16 representable values). + """ + rng = np.random.RandomState(seed) + + assert mha_type in ("mha", "mqa", "gqa") + + # --- sequence lengths (identical for Q and K) --- + lens = rng.randint(min_len, max_len + 1, size=(batch_size,)).astype(np.int64) + cu_seqlens = np.zeros(batch_size + 1, dtype=np.int32) + cu_seqlens[1:] = np.cumsum(lens) + total = int(cu_seqlens[-1]) + + # --- head counts --- + if mha_type == "gqa": + H, H_kv = 3 * n_heads, n_heads + elif mha_type == "mha": + H = H_kv = n_heads + else: # mqa + H, H_kv = n_heads, 1 + + d_head_v = 128 if d_head == 192 else d_head + + # --- generate random data in float32, then quantise to bf16 --- + def _randn_bf16(*shape): + x = rng.standard_normal(shape).astype(np.float32) + # round-trip through torch bf16 to get exact bf16 bit patterns + t = torch.from_numpy(x).to(torch.bfloat16 if dtype_str == "bfloat16" else torch.float16) + return t.float().numpy() + + q_np = _randn_bf16(total, H, d_head) + k_np = _randn_bf16(total, H_kv, d_head) + v_np = _randn_bf16(total, H_kv, d_head_v) + + # grad_out shares the shape of the output: (total_q, H, d_head_v) + grad_np = _randn_bf16(total, H, d_head_v) + + return q_np, k_np, v_np, cu_seqlens, total, grad_np + + +# --------------------------------------------------------------------------- +# Framework conversion helpers +# --------------------------------------------------------------------------- + +_TORCH_DTYPE = {"bfloat16": torch.bfloat16, "float16": torch.float16} +_PADDLE_DTYPE = {"bfloat16": paddle.bfloat16, "float16": paddle.float16} + + +def _to_torch(arr: np.ndarray, dtype_str: str, requires_grad: bool = True): + t = torch.from_numpy(arr.copy()).cuda() + t = t.to(_TORCH_DTYPE[dtype_str]) + t.requires_grad_(requires_grad) + return t + + +def _to_paddle(arr: np.ndarray, dtype_str: str, stop_gradient: bool = False): + t = paddle.to_tensor(arr.copy(), place=paddle.CUDAPlace(0)) + t = t.cast(_PADDLE_DTYPE[dtype_str]) + t.stop_gradient = stop_gradient + return t + + +# --------------------------------------------------------------------------- +# Torch reference (F.scaled_dot_product_attention, per-batch loop) +# --------------------------------------------------------------------------- + +def torch_flash_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + total_q: int, + total_k: int, + softmax_scale: Optional[float] = None, + causal: bool = False, +): + H = q.shape[1] + H_kv = k.shape[1] + B = cu_seqlens_q.shape[0] - 1 + + hcseq_q = cu_seqlens_q.cpu() + hcseq_k = cu_seqlens_k.cpu() + + outs = [] + for b in range(B): + qs, qe = int(hcseq_q[b]), int(hcseq_q[b + 1]) + ks, ke = int(hcseq_k[b]), int(hcseq_k[b + 1]) + + qb = q[qs:qe].permute(1, 0, 2).unsqueeze(0) # (1, H, Sq, d) + kb = k[ks:ke].permute(1, 0, 2).unsqueeze(0) # (1, Hkv, Sk, d) + vb = v[ks:ke].permute(1, 0, 2).unsqueeze(0) # (1, Hkv, Sk, dv) + + ob = F_torch.scaled_dot_product_attention( + qb, kb, vb, + attn_mask=None, + dropout_p=0.0, + is_causal=causal, + scale=softmax_scale, + enable_gqa=(H_kv != H), + ) + ob = ob.squeeze(0).permute(1, 0, 2).contiguous() + outs.append(ob) + + return torch.cat(outs, dim=0) + + +# --------------------------------------------------------------------------- +# Comparison logic +# --------------------------------------------------------------------------- + +def check_cross_framework( + q_np, k_np, v_np, cu_seqlens_np, total, grad_np, + dtype_str: str = "bfloat16", + softmax_scale: Optional[float] = None, + causal: bool = True, + atol: float = 3e-2, + rtol: float = 3e-2, +): + # ---- Torch side ---- + q_t = _to_torch(q_np, dtype_str) + k_t = _to_torch(k_np, dtype_str) + v_t = _to_torch(v_np, dtype_str) + cu_t = torch.from_numpy(cu_seqlens_np.copy()).cuda().to(torch.int32) + grad_t = _to_torch(grad_np, dtype_str, requires_grad=False) + + out_ref = torch_flash_ref( + q_t, k_t, v_t, + cu_seqlens_q=cu_t, cu_seqlens_k=cu_t, + total_q=total, total_k=total, + softmax_scale=softmax_scale, + causal=causal, + ) + + out_ref.backward(grad_t) + dq_ref = q_t.grad.float().cpu().numpy() + dk_ref = k_t.grad.float().cpu().numpy() + dv_ref = v_t.grad.float().cpu().numpy() + out_ref_np = out_ref.float().detach().cpu().numpy() + + # ---- Paddle side ---- + q_p = _to_paddle(q_np, dtype_str, stop_gradient=False) + k_p = _to_paddle(k_np, dtype_str, stop_gradient=False) + v_p = _to_paddle(v_np, dtype_str, stop_gradient=False) + cu_p = paddle.to_tensor(cu_seqlens_np.copy(), place=paddle.CUDAPlace(0)).cast(paddle.int32) + grad_p = _to_paddle(grad_np, dtype_str, stop_gradient=True) + + scale = (1.0 / q_np.shape[-1] ** 0.5) if softmax_scale is None else softmax_scale + + out_fa, _ = flash_attn_varlen_func( + q_p, k_p, v_p, + cu_seqlens_q=cu_p, + cu_seqlens_k=cu_p, + softmax_scale=scale, + causal=causal, + window_size=(None, None), + softcap=0.0, + return_lse=False, + ) + + out_fa.backward(grad_p) + dq_fa = q_p.grad.cast(paddle.float32).numpy() + dk_fa = k_p.grad.cast(paddle.float32).numpy() + dv_fa = v_p.grad.cast(paddle.float32).numpy() + out_fa_np = out_fa.cast(paddle.float32).numpy() + + # ---- Compare ---- + def _check(name, a, b): + diff = np.abs(a - b) + max_diff = diff.max() + ok = np.allclose(a, b, atol=atol, rtol=rtol) + if not ok: + idx = np.unravel_index(diff.argmax(), diff.shape) + print(f" {name} FAIL: max_diff={max_diff:.4e} " + f"idx={idx} fa={a[idx]:.6f} ref={b[idx]:.6f}") + else: + print(f" {name} OK: max_diff={max_diff:.4e}") + return ok + + print("---- Forward ----") + ok_fwd = _check("Out", out_fa_np, out_ref_np) + + print("---- Backward ----") + ok_dq = _check("dQ", dq_fa, dq_ref) + ok_dk = _check("dK", dk_fa, dk_ref) + ok_dv = _check("dV", dv_fa, dv_ref) + + return ok_fwd and ok_dq and ok_dk and ok_dv + + +# --------------------------------------------------------------------------- +# Pytest parametrised test +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("B", TEST_B) +@pytest.mark.parametrize("H", TEST_H) +@pytest.mark.parametrize("D", TEST_D) +@pytest.mark.parametrize("min_seq_len", TEST_MIN_SEQ_LEN) +@pytest.mark.parametrize("max_seq_len", TEST_MAX_SEQ_LEN) +@pytest.mark.parametrize("causal", TEST_CAUSAL) +@pytest.mark.parametrize("softmax_scale", TEST_SOFTMAX_SCALE) +@pytest.mark.parametrize("dtype_str", TEST_DTYPE) +@pytest.mark.parametrize("mha_type", TEST_MHA_TYPE) +def test_varlen_cross( + B, H, D, min_seq_len, max_seq_len, + causal, softmax_scale, dtype_str, mha_type, +): + if min_seq_len > max_seq_len: + pytest.skip("min_seq_len > max_seq_len") + + q_np, k_np, v_np, cu_seqlens_np, total, grad_np = generate_varlen_data( + batch_size=B, n_heads=H, d_head=D, + min_len=min_seq_len, max_len=max_seq_len, + mha_type=mha_type, dtype_str=dtype_str, + ) + + ok = check_cross_framework( + q_np, k_np, v_np, cu_seqlens_np, total, grad_np, + dtype_str=dtype_str, + softmax_scale=softmax_scale, + causal=causal, + ) + assert ok + + +# --------------------------------------------------------------------------- +# Quick standalone smoke test +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + print("=== Smoke test: B=7, H=6, D=256, mqa, causal, bf16 ===\n") + q_np, k_np, v_np, cu, total, grad_np = generate_varlen_data( + batch_size=7, n_heads=6, d_head=256, + min_len=1, max_len=2048, + mha_type="mqa", dtype_str="bfloat16", + ) + ok = check_cross_framework( + q_np, k_np, v_np, cu, total, grad_np, + dtype_str="bfloat16", + softmax_scale=None, + causal=True, + ) + print(f"\nResult: {'PASS' if ok else 'FAIL'}")