Thanks for the debugging. Could you disable autocast for self.dec_t and check the output range? This could narrow down, if the decoder outputs are indeed overflowing or if the NaNs are created in a different operation.
I’m also not sure how the quantization works in this model. Could you explain it a bit how it would quantize the values for FP32?