…does K also have to scale? I’d care a lot about wall-clock and peak memory, especially for local/on-device inference.

you can scale it based on the number of modes in your distribution, and keep it low when its not a big issue (e.g. 2/3 worked pretty well). also, this doesn't affect inference at all so on device can still be fast.