Enable MXFP8 support in TE JAX integration#424

Merged

wangye805 merged 37 commits intodevfrom

Jan 30, 2026

Commits on Dec 9, 2025

[CI] Skipped test_gpt_full_activation_recompute tests for gfx950
VeeraRajasekhar
committed
[CI] Skipped unsupported test_basic_linear_quantized tests on gfx950
VeeraRajasekhar
committed

Commits on Dec 11, 2025

[CI] Fixed test_numerics, test_norms, test_fused_optimizer failures for gfx950 ci enablement
VeeraRajasekhar
committed

Commits on Dec 12, 2025

[CI] Disabled gfx950 support until FP8 GEMM layout coverage is verified with hipblaslt
VeeraRajasekhar
committed
[CI] [gfx950] Disable cudaGraph for gemmm and grouped-gemm
VeeraRajasekhar
committed

Commits on Dec 15, 2025

Addressed reviews
VeeraRajasekhar
committed

Commits on Dec 16, 2025

Commits on Dec 18, 2025

Commits on Dec 20, 2025

Guard gfx950 TN skip by ROCm version and adjust MXFP8 Dq test size
VeeraRajasekhar
committed

Commits on Jan 6, 2026

Removed ROCM7.2 guards
VeeraRajasekhar
committed

Commits on Jan 7, 2026

Reverted ROCM7.2 guards
VeeraRajasekhar
committed

Commits on Jan 8, 2026

Merge remote-tracking branch 'origin/veergopu_gfx950_ci' into zain/jax-mxfp8
Micky774
committed

Commits on Jan 12, 2026

Corrected Normalization scale_inv padding removal
Micky774
committed

Commits on Jan 16, 2026

Updated cast behavior, added safeguards around MXFP8 GEMM
Micky774
committed

Commits on Jan 19, 2026

Commits on Jan 21, 2026

Commits on Jan 23, 2026

Commits on Jan 29, 2026