Make RoPE freq calculation broadcast for per head (#2353)
Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/2353 This is a workaround, may not be even worth landing, to avoid broadcasting semantics in the mul op and for that matter any binary op. Current implementation of oiptimized ops doesnt handle broadcasting and falls back to portable op implementation. This diff also fixes an issue where (as seen in llama) two tensors of binary op are not broadcasting, but they have different # of dims, which results in invocation of unoptimized path. e.g. a = [1, 1, 2048], b = [2048], out = [1, 1, 2048]. In llama case this is optimized path when generating one token at a time. Not so during pre-fill Making optimized op handle broadcasting, and support vectorization, is not hard, but may take some time. ghstack-source-id: 219444233 exported-using-ghexport Reviewed By: digantdesai, kirklandsign Differential Revision: D54766067 fbshipit-source-id: 0b7318959994b93388832940a98e25b9cc360978
K
Kimish Patel committed
08733f0c7239c2219c06a34f9844cd62f25a49b6
Parent: 9c20929
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 3/21/2024, 12:36:46 AM