From 56dff8d18563fa72dbd4972a562be084fc877e49 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 23:55:16 +0000 Subject: [PATCH] fix: W_gate is (H, E) but F.linear expects (E, H), transpose before linear --- dsv4/kernels/router/dense_router_decode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dsv4/kernels/router/dense_router_decode.py b/dsv4/kernels/router/dense_router_decode.py index dda5ae88..606a2f74 100644 --- a/dsv4/kernels/router/dense_router_decode.py +++ b/dsv4/kernels/router/dense_router_decode.py @@ -49,7 +49,7 @@ def _run_prefill_path( out_weights, out_ids, ): """GEMM via torch.nn.functional.linear, then fused activation + top-k.""" - logits = torch.nn.functional.linear(hidden_states.float(), W_gate.float()) + logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float()) from dsv4.kernels.router._activation_topk import run_fused_activation_topk run_fused_activation_topk( logits, e_bias, routed_scaling_factor, top_k,