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 |
||
|---|---|---|
| .. | ||
| ace | ||
| anima | ||
| audio | ||
| aura | ||
| boogu | ||
| cascade | ||
| chroma | ||
| chroma_radiance | ||
| cogvideo | ||
| cosmos | ||
| depth_anything_3 | ||
| ernie | ||
| flux | ||
| genmo | ||
| hidream | ||
| hidream_o1 | ||
| hunyuan3d | ||
| hunyuan3dv2_1 | ||
| hunyuan_video | ||
| hydit | ||
| ideogram4 | ||
| kandinsky5 | ||
| krea2 | ||
| lens | ||
| lightricks | ||
| lumina | ||
| mmaudio/vae | ||
| models | ||
| modules | ||
| moge | ||
| omnigen | ||
| pixart | ||
| pixeldit | ||
| qwen_image | ||
| rt_detr | ||
| sam3 | ||
| supir | ||
| triposplat | ||
| wan | ||
| colormap.py | ||
| common_dit.py | ||
| util.py | ||