Just don't use jit in generation and it would be fine. Of course there is some performance penalty but in my experience jit is oversold and the difference is so

Just don't use jit in generation and it would be fine. Of course there is some performance penalty but in my experience jit is oversold and the difference is something like ~10-30%. Also in any case to get optimized code you need flash attention and many other tricks.