From 33004dcbf4ad3488aab8d527e40caa52a5dc0828 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 09:22:27 +0000 Subject: [PATCH] Fix expert weight broadcasting (wt.item() for scalar multiply) --- single_shot_inference.py | 2 +- tests/test_residual_diagnostic.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index eb807f10..18c9b086 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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 ---- diff --git a/tests/test_residual_diagnostic.py b/tests/test_residual_diagnostic.py index 5c42cc95..50edcdb4 100644 --- a/tests/test_residual_diagnostic.py +++ b/tests/test_residual_diagnostic.py @@ -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