DispatchKeySet perf improvements (#72828)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72828
Reland of D34034847 (https://github.com/pytorch/pytorch/commit/8aa3620d73d9afcdbcf3d691cdd45ae05ec5e7d5)
ghstack-source-id: 152161453
Test Plan: confirm that Milan tests are passing
Reviewed By: ezyang, albanD
Differential Revision: D34227615
fbshipit-source-id: c7695e16dba3076e8ab9df8654327c5d57e92c77
(cherry picked from commit 940717db1551b799964894e0bb97757ecae14235)