xla
1ccd6a64 - [Pallas] Add a bfloat16 flash attention test case (#6748)

Commit

1 year ago

[Pallas] Add a bfloat16 flash attention test case (#6748) Summary: Add a bfloat16 flash attention test case. Test Plan: python test/test_pallas.py

References

Author

alanwaketan

Parents