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

[None][chore] enable EPLB for DEEPGEMM (#10617)

Browse Source 
Signed-off-by: xxi <xxi@nvidia.com>
This commit is contained in:
xxi 2026-01-14 18:28:08 +08:00 committed by GitHub
parent 272688c663
commit d8862505b9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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