Use optimized argsort and topk of aten native op in MoE or other requiring sort.

[voidbag]

refer to the draft PR 2343
pytorch/pytorch#170978 (comment)

📊 Detailed torch.sort(indices_dtype) Performance Benchmark Results & Latency Analysis

✅ Status & Todo

• Integrate _indices_dtype_by_sort_size logic into torchtitan.distributed.deepep.
• Update TokenReorderer in torchtitan.models.moe.
• Work In Progress pytorch/pytorch#170978
• optional When sort is merged, topk can be modified to support indices_dtype in out-variant op.
• optional Support general integer indexing to enjoy fully indices dtype feature of torch.sort.