GitLab

[Kernel] Support MoE Fp8 Checkpoints for Mixtral (Static Weights with Dynamic/Static Activations) (#4527)

Follow on to #4332 to enable FP8 checkpoint loading for Mixtral and supersedes #4436.

This PR enables the following checkpoint loading features for Mixtral:

Supports loading fp8 checkpoints for Mixtral, such as this "nm-testing/Mixtral-8x7B-Instruct-v0.1-FP8" test model
Supports static or dynamic activation quantization with static weight quantization (all per tensor)
Supports different scales for each expert weight
Supports Fp8 in QKV layer
Notes:

The Expert Gate/Router always runs at half / full precision for now.
If there are different weight scales between QKV layer (for separate QKV weights), they are re-quantized using layer.weight_scale.max() so we can have a single gemm for performance.

Co-authored-by: LiuXiaoxuanPKU <llilyliupku@gmail.com>

Co-authored-by: simon-mo <simon.mo@hey.com>

Co-authored-by: SangBin Cho <rkooo567@gmail.com>

Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>

Co-authored-by: Cade Daniel <edacih@gmail.com>

[Bug fix][Core] assert num_new_tokens == 1 fails when SamplingParams.n is not 1 and max_tokens is large & Add tests for preemption (#4451)

Remove the device="cuda" declarations in mixtral as promised in #4343