Unverified Commit 121f92c5 authored by Yuan Luo's avatar Yuan Luo Committed by GitHub
Browse files

Add main for merge state tests (#6492)


Co-authored-by: default avatarluoyuan.luo <luoyuan.luo@antgroup.com>
parent 3bde1010
......@@ -136,3 +136,7 @@ def test_merge_state(seq_len, num_heads, head_dim):
assert torch.allclose(v_merged, v_merged_std, atol=1e-2)
assert torch.allclose(s_merged, s_merged_std, atol=1e-2)
if __name__ == "__main__":
pytest.main([__file__])
......@@ -394,3 +394,7 @@ def test_merge_attn_states(
len(NUM_BATCH_TOKENS) * len(HEAD_SIZES) * len(NUM_QUERY_HEADS) * len(DTYPES)
):
generate_markdown_table()
if __name__ == "__main__":
pytest.main([__file__])
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment