[RL] Support FlashInfer TRT-LLM NVFP4 MoE in the RL weight checker (#31085)

Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>
2026-07-24 16:01:54 -07:00
parent 0a212c6119
commit 14d6e1d3b1
1 changed files with 7 additions and 1 deletions
@@ -161,7 +161,13 @@ def select_comparable_weight(quant_method) -> Optional[type]:
        and not quant_method.use_mxfp8
    ):
        return Fp8BlockComparable
-    if isinstance(quant_method, (ModelOptFp4LinearMethod, ModelOptNvFp4FusedMoEMethod)):
+    if isinstance(quant_method, ModelOptNvFp4FusedMoEMethod):
+        if getattr(quant_method, "enable_flashinfer_trtllm_moe", False):
+            return None
+        raise NotImplementedError(
+            f"weight checker has no ComparableWeight for {type(quant_method).__name__}"
+        )
+    if isinstance(quant_method, ModelOptFp4LinearMethod):
        raise NotImplementedError(
            f"weight checker has no ComparableWeight for {type(quant_method).__name__}"
        )