dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(CRITICAL): corex_gdn Conv1d groups=kd — was crashing on first prefill

Browse Source 
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.
This commit is contained in:
project6-dev
2026-08-11 02:58:39 +00:00
parent d1c5e992aa
commit b25fc53e5c
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
ex_engine/python/corex_gdn.py
View File

@@ -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