xla
custom_kernel: fix shape mismatch by sharding segment_ids in flash attn.
#8333

Open

custom_kernel: fix shape mismatch by sharding segment_ids in flash attn. #8333

dudulightricks wants to merge 2 commits into pytorch:master from dudulightricks:bug-fix/shard-segment-ids-flash-attention

dudulightricks marked this pull request as draft 1 year ago

miladm commented on 2024-11-14

dudulightricks force pushed from 4ff4357b to 1f1a3bd8 1 year ago

dudulightricks force pushed from 1f1a3bd8 to 0c791645 1 year ago

dudulightricks force pushed from 75caebb1 to 5b29a99d 1 year ago

dudulightricks force pushed from fa4cfba7 to 0e668fe9 1 year ago

dudulightricks force pushed from 0e668fe9 to 804da333 1 year ago

dudulightricks force pushed from 804da333 to b1e3ffa9 1 year ago

dudulightricks force pushed from b1e3ffa9 to c3ef029a 1 year ago

dudulightricks force pushed from 8d612bfa to 8079ed7e 1 year ago

dudulightricks force pushed from 8079ed7e to e3c13a0f 1 year ago

dudulightricks force pushed from e3c13a0f to 5c13ea29 1 year ago

dudulightricks force pushed from 42aded00 to 9dedcaa4 1 year ago

dudulightricks force pushed from 9dedcaa4 to 02cd100d 1 year ago

dudulightricks force pushed from 02cd100d to e7573d33 1 year ago

dudulightricks force pushed from e7573d33 to 9b3891b9 1 year ago

dudulightricks force pushed from 9b3891b9 to 3697955e 1 year ago

dudulightricks force pushed from 3697955e to d3e06bc0 1 year ago

dudulightricks force pushed from d3e06bc0 to e07f8034 1 year ago

dudulightricks force pushed from e07f8034 to cdd29e7f 1 year ago

dudulightricks force pushed from cdd29e7f to b231efae 1 year ago

dudulightricks force pushed from b231efae to a3fb4e15 1 year ago

dudulightricks force pushed from 68134fc1 to 32ba86c1 1 year ago

dudulightricks force pushed from 32ba86c1 to 8e69c26e 1 year ago

dudulightricks force pushed from 8e69c26e to d66d0687 1 year ago

dudulightricks force pushed from d66d0687 to d8d089be 1 year ago

dudulightricks force pushed from d8d089be to 08410f95 1 year ago

dudulightricks force pushed from 08410f95 to 4ba90671 1 year ago

dudulightricks force pushed from 4ba90671 to 06089004 1 year ago

miladm added pallas

custom_kernel: fix shape mismatch by sharding segment_ids in flash attn.

8b096010

test proposel for segment ids that fails.

b5d1b8f6

dudulightricks force pushed from 06089004 to b5d1b8f6 1 year ago

Reviewers

miladm

Assignees

No one assigned

Labels

pallas

Milestone

No milestone