MXFP8 for experts in MoE

[kailash109]

Adds support MXFP8 quantization of expert layers in MoE models

Passes through config details for mxfp8 quantization to both Megatron provider and VLLM quantization config

Keeping this as a draft PR as currently testing on 8xB200 on Qwen-30B-A3B observes noticeable perf degradation from enabling mxfp8 vs bf16:

Reproduce with uv run --isolated --with modal modal run --detach \ examples/train/megatron/modal_expert_mxfp8_benchmark.py \ --mode {bf16/fp8} --steps 10

Will update with profiling results/updates

[erictang000]

hey @kailash109, we're currently working on this in #1898 (both mxfp8 on blackwell and fp32 scaling on hopper)

there should be some relevant details there for what was needed to get things working beyond just plumbing through the configs

[kailash109]

Hi @erictang000 thanks for the response -- I haven't looked into persistent mxfp8 for expert weights, instead I was running some experiments based on this to see what kind of tps improvement I can get for Moe while keeping weights in bf16 and quantizing on the fly (likewise using vllm's online quantization from bf16 -> mxfp8 in weight syncs)

I haven't checked my profile traces yet, but I suspect a part of the perf degredation I'm seeing is due to the additional quantization overhead in the expert forward/backward, so I'll look into persistent mxfp8 for expert layers.

To my understanding, #1898 is doing fp8 quantization dsv3-style with fp32 scales for 128x128 blocks, but it seems I can reuse a good bit of the logic for weight syncs -- will keep you posted!

[erictang000]

got it!

i think #1898 does also include support for mxfp8 on blackwell in addition to the fp32 scaling for hopper, but i think adding on configs to quantize experts specifically makes sense as well, let me know what you find

cc: @YJHMITWEB fyi

[erictang000]

ah @kailash109 i see your point on the block size being hard coded to 128x128 (using fp8_recipe=blockwise vs mxfp8) in the linked pr even if it's using MXFP8 style power of 2 scaling rather than fp32, it would be good to make sure these PRs are overall compatible and if we make the configuration similar for both settings

[kailash109]

Hey @erictang000 , I've created draft PR #1940 which builds persistent MXFP8 on top of this PR and the FP8 E2E training PR (it uses the same blockwise_cast and iter_serialized_fp8_tensors from @YJHMITWEB to pack MXFP8-style weights for VLLM, as you mentioned the main change was providing the 1x32 block size into the function + packing the E8M0 power-of-2 scale factor into single byte)

I still have to debug perf degradations I see there compared to bf16 baseline which is why I've left it as draft, the next step is to create a generalized QuantizationStrategy which I'm designing to support arbitrary serializers for different quantization schemes as well as target modules (which will be helpful in the future for NVFP4 support)