add control over the number of SMs to be used by the kernel

[Amir-19]

added optional parameter max_sm_count to the kernel to specify the number of SMs to use in the kernel.

[nandor]

I think this is a good addition, however the parameter would be better placed on the calls to dispatch/combine themselves. This way, a single instantiation of all-to-all can be re-used in multiple contexts (prefill/decode/cuda-graph/no-cuda-graph).