vllm - 💡(How to fix) Fix [Feature]: Integrate fused `kMoEFinalizeARResidualRMSNorm` from FlashInfer [1 comments, 2 participants]

TL;DR

Extract the moe_finalize from the wrapped fused_moe op to create a new torch.compile custom pass.

Guidance

• Identify the fused_moe op in the FlashInfer 0.6.8 codebase and locate the embedded moe_finalize to determine the best approach for extraction.
• Verify that the moe_finalize can be successfully pulled out and used as a separate custom pass in torch.compile.
• Investigate the PR https://github.com/flashinfer-ai/flashinfer/pull/2982 for potential insights into the fused MoE implementation.
• Consider testing the new custom pass with a minimal example to ensure it functions as expected.

Notes

The solution assumes that the moe_finalize can be extracted and used as a custom pass without significant modifications to the existing codebase.

Recommendation

Apply workaround: Extract and refactor the moe_finalize to create a new torch.compile custom pass, as this approach allows for incremental testing and validation of the new pass without requiring a full upgrade or rewrite.