[BUG] ONNX Export fails for Qwen3.5 MTP draft model with 4-bit quantized checkpoint due to unpacked mtp.fc.weight shape mismatch

[fffiore]

Bug Description

When attempting to export an int4_awq quantized Qwen3.5-4B checkpoint to ONNX using tensorrt-edgellm-export with the --mtp flag, the export process crashes during the [MTP Draft] tracing phase with a matrix shape mismatch error.

Reproduction Steps

1. Quantize the model:

   tensorrt-edgellm-quantize llm \
       --model_dir /path/to/Qwen3.5-4B \
       --output_dir /path/to/quant/091/Qwen3.5-4B/ \
       --quantization int4_awq
   

2. Export to ONNX with --mtp:
   tensorrt-edgellm-export
   /path/to/quant/091/Qwen3.5-4B
   /path/to/onnx/091/Qwen3.5-4B/
   --mtp

3. The export crashes during MTP draft ONNX export with the following RuntimeError:

File ".../tensorrt_edgellm/models/qwen3_5/modeling_qwen3_5_mtp.py", line 230, in forward
fused_hidden_states = self.fc(torch.cat((normed_embeds, normed_hidden_states), dim=-1))
File ".../tensorrt_edgellm/models/linear.py", line 171, in forward
return F.linear(hidden_states, self.weight, bias)

RuntimeError: a and b must have same reduction dim, but got [s23*s30, 5120] X [2560, 2560].
torch.onnx._internal.exporter._errors.TorchExportError: Failed to export the model with torch.export.