transformers
955b6d35 - Restore the `MegaBlocksMoeMLP` hub kernel mapping for MXFP4 GPT-OSS (#50150)

Commit
8 days ago
Restore the `MegaBlocksMoeMLP` hub kernel mapping for MXFP4 GPT-OSS (#50150) * Restore MegaBlocksMoeMLP hub kernel mapping for MXFP4 GPT-OSS MLP #49548 replaces GptOssMLP with Mxfp4GptOssMLP, which dropped the MegaBlocksMoeMLP hub mapping, so use_kernels=True no longer kernelizes the MoE MLP of MXFP4 GPT-OSS on CPU/XPU. * Test the megablocks kernel on a real MXFP4 GPT-OSS model on CPU/XPU Replaces the decorator-only check with a slow test that loads openai/gpt-oss-20b with use_kernels=True, checks every MLP keeps MXFP4 experts and runs the MegaBlocksMoeMLP kernel, and checks generation. Limited to CPU/XPU: other devices dequantize to bf16 with use_kernels=True and keep GptOssMLP. * Use the modeling flash-attn2 kernel version in benchmark_v2's FA2 check is_fa2_or_kernel_available() pinned kernels-community/flash-attn2 to version=1, which has no XPU build, so flash_attention_2 benchmarks silently fell back to SDPA on XPU. Use FLASH_ATTN_KERNEL_VERSIONS, the same version the model loads. --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
Author
Parents
Loading