Commits · 01a504c4bdbb21aa2bfaddb8fc49c3a7f0e1fbac · OpenDAS / TransformerEngine

16 Jun, 2025 1 commit

[JAX] Grouped GEMM & Dense support MXFP8 and handle empty matrices (#1871) · 01a504c4

Hua Huang authored Jun 16, 2025



* Support MXFP8 and handle empty matrices
Signed-off-by: Hua Huang <huah@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>

01a504c4

06 Jun, 2025 1 commit

[JAX] GroupedQuantizer and GroupedScaledTensor (#1666) · 7948779c

Phuong Nguyen authored Jun 06, 2025



* refactor the multi_stream utils + implement nvte_multi_tensor_quantize in TE/Common

* implement GroupedQuantizer and grouped_quantize in jaxx

* fix logical_axes_names for transpose tensor in ScaledTensor
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Hua Huang <huah@nvidia.com>
Co-authored-by: Ming Huang <mingh@nvidia.com>

7948779c