Max calculation inside of forward creates CUDA sync

[javak87]

In the attention mechanism’s forward pass, we compute max(x_q_lens) and max(x_kv_lens) in order to call flash_attn_varlen_func. This max operation triggers a CUDA synchronization, which prevents effective use of torch.compile. Since x_q_lens and x_kv_lens are inputs (not intermediates), it’s recommended to move this CUDA sync as far up the call stack as possible (i.e., handle it at a higher level) to minimize its impact.

[tjhunter]

I wonder if we can hardcode the values. Otherwise it will cause other issues with torch.compile: if the length is different at each step, then this will trigger a model compilation at each step. This negates most of the advantages of compiling the model to begin with. We limit the number of tokens for each cell, so that could be a reasonable value.