kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-13 22:18:36 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] enable EPLB for DEEPGEMM

Browse Source 
Signed-off-by: xxi <xxi@nvidia.com>
This commit is contained in:
xxi 2026-01-13 03:26:50 +00:00
parent ba1cb6831d
commit 8d998af03b
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorrt_llm/_torch/modules/fused_moe/create_moe.py
View File

@ -138,8 +138,9 @@ def create_moe_backend(
moe_load_balancer = get_moe_load_balancer()
if moe_load_balancer is not None:
assert moe_cls in [
WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE
], "MoE Load Balance is only supported in WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE and CuteDslFusedMoE now."
WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE,
DeepGemmFusedMoE
], "MoE Load Balance is only supported in WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE and CuteDslFusedMoE, and DeepGemmFusedMoE."
if bias:
assert moe_cls in [CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE