Support Scaled matmul and scaled dot_general in jax.nn

[wenscarl]

This PR introduces two new nn APIs:

1. scaled_matmul – Performs matrix multiplication with scaling factors.
2. scaled_dot_general – Implements a generalized dot product operation with quantization, matmul and dequantization.
   Currently, only the e8m0fnu type scaling factor is supported.
   @kaixih @nouiz

[kaixih]

@IvyZX @sdasgup3 can you help review or find other reviewers?

[IvyZX]

PTAL test errors - looks like e8m0fnu is a not type that Numpy supports. I'm not sure if we should adopt this type in JAX.

Also, is it our common practice to add purely-cudnn APIs to jax.nn? Looks like in #21371 an XLA implementation is provided as an alternative for users without cudnn.

Adding maintainers who know more about JAX's lower level common practices: @hawkinsp @superbobry

[kaixih]

PTAL test errors - looks like e8m0fnu is a not type that Numpy supports. I'm not sure if we should adopt this type in JAX.

This dtype should already in the JAX (see here). The dtype is first defined in ml_dtype since 0.5.0 and imported to JAX.

Also, is it our common practice to add purely-cudnn APIs to jax.nn? Looks like in #21371 an XLA implementation is provided as an alternative for users without cudnn.

The scaled_matmul and scaled_dot_general are supposed to be general and not cudnn-specific. Yes, both of them call into the cudnn_scaled_matmul, which appears to be purely-cudnn API. But cudnn_scaled_matmul will actually fall back to XLA impl if cudnn doesn't support the case (see this XLA pass). So, that explains why we don't see the XLA impl in the python code.

[sdasgup3]

cc @reedwm

[IvyZX]

Please also make sure all CI tests pass.

[IvyZX]

Note that JAX doesn't require ml_dtypes >= 0.5.0, so jax.numpy.float8_e8m0fnu is not guaranteed to exist.

Please also avoid defining a default fp8 config as a global instance - you can always pass it in via jax.nn level API.

[IvyZX]

See CI test failures - this code only runs on GPU now.

You will need to either:

• add lowering fallback to XLA for other platforms like CPU and TPU, or
• explicitly say your implementation only works for certain platforms and prevent invalid usages: platform-check or require an explicit platform arg in the API (like this cudnn implementation), and disable tests on other platforms.