Commit 1b54822
ssjia
Update base for Update on "[ET-VK][ez] Fuse update_cache + custom_sdpa into sdpa_with_kv_cache"
SDPA used to be handled by a custom op `sdpa_with_kv_cache`, but it was eventually split (D62301837) into update_cache and custom_sdpa ops.
However, having a single fused op is useful for Vulkan since it allows more control over how the cache tensors are stored and represented. Essentially, it makes it easier to manage the cache tensors and opens up opportunities for future optimizations. This diff introduces a fusion pass that does 2 things:
1. Combine update_cache and custom_sdpa back into sdpa_with_kv_cache
2. Ensure all references to the cache_pos symint use the same node - this prevents the select_at_dim_as_symint op from being called every time it is used.
Differential Revision: [D86340339](https://our.internmc.facebook.com/intern/diff/D86340339/)
[ghstack-poisoned]1 parent 1251adb commit 1b54822
File tree
0 file changed
+0
-0
lines changed0 file changed
+0
-0
lines changed
0 commit comments