trunk/5e0379393cc64f15209ec20e20bc2eddc5ed2a42: Add autograd support to TokenSwitch dispatch and combine (#181314)

dispatch and combine are adjoint operations: backward of dispatch calls

combine, and backward of combine calls dispatch. Single dispatch/combine

public API following nn.Module/torch.matmul conventions:

• With out=(out_tokens, out_weights, out_idx): writes to caller-supplied
  
  buffers and returns them; no autograd (efficient buffer reuse).


• Without out: allocates buffers internally, returns them with full
  
  autograd support via _DispatchAutograd / _CombineAutograd.

Subclasses implement _dispatch/_combine as the raw buffer-writing

primitives and inherit both modes from the base class. topk_weights

receives no gradient (routing metadata at a different byte width).

Authored with Claude.

Pull Request resolved: #181314

Approved by: https://github.com/kapilsh

ghstack dependencies: #178712