55760cdeb5
Addresses CodeRabbit feedback on #14838: the chunking loop always allocated a temporary output buffer and copied SDPA's result into it, even in the common case (any non-MPS device, or MPS below the size threshold) where steps==1 and the loop only ever runs once. Adds a fast path that calls scaled_dot_product_attention directly and reshapes its result, matching the original pre-fix code for that case -- no extra allocation or copy. The chunked path (steps>1) is unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW |
||
|---|---|---|
| .. | ||
| diffusionmodules | ||
| distributions | ||
| encoders | ||
| attention.py | ||
| ema.py | ||
| sdpose.py | ||
| sub_quadratic_attention.py | ||
| temporal_ae.py | ||