From 1827ec57f746bc70c005bea9174f185ab1d35969 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 7 Aug 2026 23:42:02 +0800 Subject: [PATCH] gptoss: fix sharded invalids (#17450) --- extra/gemm/moe_routing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/extra/gemm/moe_routing.py b/extra/gemm/moe_routing.py index e683b83067..30ba001c3f 100644 --- a/extra/gemm/moe_routing.py +++ b/extra/gemm/moe_routing.py @@ -5,7 +5,8 @@ BLOCK_ROW = 256 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor: if isinstance(device, tuple): - return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0) + per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device) + return Tensor(per.uop.unshard(0), device=device) return Tensor.invalids(*shape, dtype=dtype, device=device) def _atomic_add(device:str) -> str: