add checks in trtllm gen fused moe module

pamelap-nvidia · pamelap-nvidia · commit 66b9f2220b7d · 2025-08-20T14:54:02.000Z
Signed-off-by: Pamela &lt;179191831+pamelap-nvidia@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
@@ -68,6 +68,9 @@ def __init__(
             weight_loading_mode=weight_loading_mode,
         )
 
+        sm_version = get_sm_version()
+        assert sm_version < 120, "TRTLLMGenFusedMoE does not support SM120 and SM121."
+
         assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE."
 
         self.num_slots = self.num_experts

Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,9 @@ def __init__(`
`68`	`68`	`weight_loading_mode=weight_loading_mode,`
`69`	`69`	`)`
`70`	`70`
	`71`	`+ sm_version = get_sm_version()`
	`72`	`+ assert sm_version < 120, "TRTLLMGenFusedMoE does not support SM120 and SM121."`
	`73`	`+`
`71`	`74`	`assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE."`
`72`	`75`
`73`	`76`	`self.num_slots = self.num_experts`