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

Fix expert weight broadcasting (wt.item() for scalar multiply)

Browse Source
This commit is contained in:
biondizzle
2026-05-31 09:22:27 +00:00
parent 1434b35971
commit 33004dcbf4
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
single_shot_inference.py
View File

@@ -572,7 +572,7 @@ def moe_forward(x, w, li, cfg, token_id, device):
# Weighted combine + scaling
routed_out = torch.zeros_like(x)
for i, (out, wt) in enumerate(zip(expert_outputs, expert_weights)):
routed_out = routed_out + (out.float() * wt).bfloat16()
routed_out = routed_out + (out.float() * wt.item()).bfloat16()
routed_out = (routed_out.float() * routed_scaling).bfloat16()
# ---- Shared expert ----

2
tests/test_residual_diagnostic.py
View File

@@ -215,7 +215,7 @@ def main():
routed_out = torch.zeros_like(x_ffn_normed)
for i, (out, wt) in enumerate(zip(expert_outputs, expert_weights)):
routed_out = routed_out + (out.float() * wt).bfloat16()
routed_out = routed_out + (out.float() * wt.item()).bfloat16()
routed_out = (routed_out.float() * routed_scaling).bfloat16()
# Shared expert