[torch.compile] Unwrap fused_marlin_moe custom op (#26739)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
This commit is contained in:
Varun Sundar Rabindranath
2025-10-13 22:22:16 -04:00
committed by GitHub
parent 8a0af6a561
commit 8ae169286f
10 changed files with 22 additions and 52 deletions

View File

@@ -21,7 +21,10 @@ from vllm.model_executor.layers.fused_moe.config import (
mxfp4_w4a16_moe_quant_config,
ocp_mx_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import MarlinExperts
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
MarlinExperts,
fused_marlin_moe,
)
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
OAITritonExperts,
)
@@ -947,7 +950,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
e_score_correction_bias=e_score_correction_bias,
)
return torch.ops.vllm.fused_marlin_moe(
return fused_marlin_moe(
x,
layer.w13_weight,
layer.w2_weight,