Use TransformerEngine for MXFP8 quantization

[zianglih]

@HumansAnd

Summary

• replace the FlashInfer/Triton MXFP8 quantizer selection with one shared TransformerEngine adapter
• reuse the TE path for both Megatron-to-HF live weight export and offline HF checkpoint conversion
• remove the obsolete SGLang mxfp8_group_quantize re-export and its CPU import stub

Implementation

The shared helper lazily constructs the TransformerEngine 2.17 MXFP8Quantizer in rowwise-only mode. It flattens leading dimensions, pads rows to TE's 32-row requirement, and crops the quantized payload and UE8M0 scales back to the existing compact, unswizzled checkpoint contract.

Existing layer selection, BF16 carve-outs, source block-FP8 handling, tensor names, and output layouts are unchanged.

Validation

• pre-commit run --all-files
• B200 devbox using radixark/miles:dev-202607290235
  • python3 docker/verify_transformer_engine.py transformer_engine_cu13
  • TransformerEngine version: 2.17.0
  • CUDA_VISIBLE_DEVICES=0 python3 -m pytest tests/fast-gpu/test_mxfp8_quantizer.py -xvs (108 passed)
  • CUDA_VISIBLE_DEVICES=0 python3 -m pytest tests/fast/backends/megatron_utils/test_hf_weight_iterator_direct.py -xvs (6 passed)
  • DeepSeek V3.2 5-layer MXFP8 E2E: PYTHONPATH=/root/miles python3 tests/e2e/megatron/model_scripts/test_deepseek_v32_5layer_mxfp8.py (passed)
    • completed a fresh FP8 to BF16 to MXFP8 conversion with TransformerEngine 2.17.0
    • quantized weight comparison used allow_quant_error=True; /weights_checker returned HTTP 200
    • rollouts 0, 1, and 2 completed with weight versions 1, 2, and 3 and mixed_version_ratio=0.0
    • Ray job raysubmit_UP189w1SijVshDUf succeeded

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.