From e873e5f27b1ac8711e989708e1381939e62a19c1 Mon Sep 17 00:00:00 2001 From: dylan Date: Sat, 15 Aug 2026 13:14:11 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20eliminate=208x=20CUDA=20sync=20in=20MoE?= =?UTF-8?q?=20decode=20=E2=80=94=20tolist()=20once=20instead=20of=20.item(?= =?UTF-8?q?)=20per=20expert?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ex_engine/moe/naive_batched_experts.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ex_engine/moe/naive_batched_experts.py b/ex_engine/moe/naive_batched_experts.py index ecc4b11e..f1656312 100644 --- a/ex_engine/moe/naive_batched_experts.py +++ b/ex_engine/moe/naive_batched_experts.py @@ -65,11 +65,11 @@ def naive_batched_moe_forward( # # For decode, each expert sees exactly 1 token. # expert ids are in topk_ids[0] (shape: top_k,) - eids = topk_ids[0] # (top_k,) - ws = topk_weights[0] # (top_k,) + eids = topk_ids[0].tolist() # (top_k,) → CPU list, ONE sync + ws = topk_weights[0] # (top_k,) stays on GPU for i in range(top_k): - eid = eids[i].item() + eid = eids[i] # FC1: (1, H) @ (H, 2*I) → (1, 2*I) # w13[eid] is (2*I, H), .transpose(0, 1) is (H, 2*I) — VIEW, zero copy