I did a quick simulation-based test and it works for me. The only changes required are in cw305_top.v: pKEY_WIDTH=256 and aes_core.size_i=2.
If you changed pKEY_WIDTH in cw305_reg_aes.v only, it’ll get overwritten by the 128 that gets passed down from cw305_top.v.