biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: call moe_runner.run() and se_runner.run() (not __call__)

Browse Source
This commit is contained in:
biondizzle
2026-06-01 01:14:38 +00:00
parent 03c45d4bfb
commit 91b3929605
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
single_shot_inference.py
View File

@@ -380,7 +380,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# =====================================================================
def moe_forward(x, li, moe_runner, se_runner, router, token_id):
topk_w, topk_ids = router(x, token_ids=token_id)
routed_out = moe_runner(x, topk_w, topk_ids); shared_out = se_runner(x)
routed_out = moe_runner.run(x, topk_w, topk_ids)
shared_out = se_runner.run(x)
return routed_out + shared_out
# =====================================================================