[Hotfix] Fix router gemm on sm103 (#22134)

This commit is contained in:
Baizhou Zhang
2026-04-05 09:33:14 -07:00
committed by GitHub
parent f6c9072e42
commit c5fa364b80
+1 -1
View File
@@ -326,7 +326,7 @@ class MoEGate(nn.Module):
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
and _device_sm >= 90
):
if _device_sm >= 100 and self.weight.shape[0] == 256:
if _device_sm == 100 and self.weight.shape[0] == 256:
# router gemm output float32
logits = torch.empty(
hidden_states.shape[0],