Refactor Megablox GMM kernels for quantization

* Rename gmm.py to backend.py to avoid the name conflict with gmm function.
* Make gmm() function a wrapper around the actual custom_vjp object.
* Move quantization of parameters from backend.py to ops.py for better consistency.
* In bwd pass, apply the scales from residuals to the gradient to avoid double quantization.

PiperOrigin-RevId: 814436757

Author: liudangyi