Chunked full-attention SDPA for long key sequences#3307

Open

Thump604 wants to merge 7 commits intoml-explore:mainfrom

Thump604:feat/chunked-sdpa

Commits on Mar 24, 2026

test: add logsumexp output correctness tests for fused SDPA
Thump604
committed
feat: add output_logsumexp function constant to steel_attention kernel
Thump604
committed
feat: enable logsumexp output in fused SDPA dispatch
Thump604
committed
test: add chunked SDPA correctness and edge case tests
Thump604
committed
feat: add chunked SDPA reduction kernel for logsumexp-weighted merging
Thump604
committed
feat: add chunked SDPA dispatch for large key sequences
Thump604
committed
test: add integration tests at 128K/256K context to prove GPU watchdog fix
Thump604
committed