From b25fc53e5cbc5a61b01e9b7b9eb6edd276cfb2b5 Mon Sep 17 00:00:00 2001 From: project6-dev Date: Tue, 11 Aug 2026 02:58:39 +0000 Subject: [PATCH] =?UTF-8?q?fix(CRITICAL):=20corex=5Fgdn=20Conv1d=20groups?= =?UTF-8?q?=3Dkd=20=E2=80=94=20was=20crashing=20on=20first=20prefill?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Error: 'Given groups=1, weight [1,1,4], expected input [1,128,4099] to have 1 channels but got 128' Root cause: kh_pad is (kd, N+pad) = (128, 4099), but weight was (1, 1, 4) with groups=1. Conv1d requires in_channels == input_channels/groups, so 1 != 128/1. Fix: expand weight to (kd, 1, conv_kernel_size) and use groups=kd for depthwise conv. This matches the pattern in qwen3_5.py:212 (_causal_conv1d_fwd) which uses groups=channels. This was the cause of 'evaluation failed' — GDN crash on first request killed the engine. --- ex_engine/python/corex_gdn.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/ex_engine/python/corex_gdn.py b/ex_engine/python/corex_gdn.py index a8d143b1..cd679674 100644 --- a/ex_engine/python/corex_gdn.py +++ b/ex_engine/python/corex_gdn.py @@ -92,15 +92,18 @@ class CoreXGDN: # Prefill: apply conv1d directly on sequence k_conv = k.transpose(0, 1).unsqueeze(0) # (1, nk, N, kd) # Reshape for grouped conv: (1, nk, N, kd) -> (nk, 1, N) per head, apply conv + # Depthwise conv1d per head, matching qwen3_5.py _causal_conv1d_fwd pattern + # conv1d_weight: (nk, 1, conv_kernel_size) k_out = [] for h in range(nk): kh = k_conv[0, h] # (N, kd) - # Pad and conv each dim independently? No — conv is on seq dim kh_t = kh.t() # (kd, N) - kh_pad = F.pad(kh_t, (self.conv_kernel_size - 1, 0)) # causal pad + kh_pad = F.pad(kh_t, (self.conv_kernel_size - 1, 0)) # causal pad: (kd, N+pad) + # Depthwise: each of kd channels gets its own conv with same weight w = conv1d_weight[h] # (1, conv_kernel_size) - kh_conv = F.conv1d(kh_pad.unsqueeze(0), w.unsqueeze(0).float(), - groups=1).squeeze(0)[:, :num_tokens] + w_expand = w.expand(kd, -1).unsqueeze(1).float() # (kd, 1, conv_kernel_size) + kh_conv = F.conv1d(kh_pad.unsqueeze(0), w_expand, + groups=kd).squeeze(0)[:, :num_tokens] # (kd, N) k_out.append(kh_conv.t()) # (N, kd) k = torch.stack(k_out, dim=1).to(hidden_states.dtype) # (N, nk, kd) # Update conv_state for decode