First warmup with the same input/output as benchmark

[IlyasMoutawwakil]

fixes an issue with torch.compile capturing a small static kv-cache tensor during warmup, that's used afterwards with longer token generation queries, resulting in device-side errors (index errors).
also adds torch_compile_target which can be either model or forward.

@ydshieh