Skip to content

Support d != dv test, add varlen test. #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
baoqiwen wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Support d != dv test, add varlen test. #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 38 additions & 25 deletions benchmark_fa4_mask_mod.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,7 @@ def test_mask(
HKV,
S,
D,
DV,
dtype,
skip_correctness: bool = False,
print_mask: bool = True,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -842,16 +844,21 @@ 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

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:
Expand All @@ -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:
Expand Down Expand Up @@ -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()

Expand Down
Loading