xla
Fix a bug in flash attention where kv_seq_len should divide block_k_major.
#8671

Merged

Fix a bug in flash attention where kv_seq_len should divide block_k_major. #8671

zpcore merged 8 commits into pytorch:master from zhangp365:master

Fix a bug in flash attention where kv_seq_len should divide block_k_m…

45f6e963

qihqi requested a review from

tengyifei 270 days ago

qihqi requested a review from

zpcore 270 days ago

qihqi requested a review from

qihqi 270 days ago

zpcore commented on 2025-02-06

add test code for flash attention and adjust previous code.

d4c84198

fixed a bug in test code.

cb583e66

adjust the code format in test_pallas.py and test_pallas_spmd.py by y…

5d583a6e

adjust the code format in test_pallas.py and test_pallas_spmd.py by y…

0a639e87

zpcore commented on 2025-02-07

fix a padding bug in custom_kernel.py

cb2157f0

fixed a padding softmax bug in custom_kernel.py

41ac7acc

move the padding position in custom_kernel.py

4f0532dd

zpcore approved these changes on 2025-02-10

zpcore merged cff9f4e0 into master 265 days ago

Reviewers

zpcore

tengyifei

qihqi

Assignees

No one assigned

Labels

None yet

Milestone

No milestone