fix(moe): Handle dtype mismatch in torch._grouped_mm with autocast (#43839)
* fix(moe): Handle dtype mismatch in grouped_mm with autocast
torch._grouped_mm is not autocast-enabled, so when using torch.autocast
with MoE models like Phi-tiny-MoE, the input tensor may have a different
dtype than the weights, causing RuntimeError.
Add explicit dtype casting in _grouped_linear to ensure input and weight
tensors have the same dtype before calling torch._grouped_mm.
Fixes #43828
* refactor: use autocast(enabled=False) instead of manual dtype cast
Replace manual input.to(weight.dtype) cast with torch.amp.autocast(enabled=False)
wrapper around torch._grouped_mm calls. This is cleaner since it addresses the
root cause (grouped_mm not being autocast-aware) rather than guessing the target
dtype.
Ref: https://github.com/pytorch/pytorch/issues/174763