Fix memory calculation causing --compile to error

[ajtejankar]

The exact reason of the error is not clear to me, but there's a simpler way to calculate the free memory which is implemented here. Here's what I've understood so far.

1. When compile is enabled, the code path that estimates cuda graph memory causes PyTorch to release temporary memory which increases the amount of free memory available
2. This doesn't happen when compile is not enabled which results in lesser free memory
3. My guess is that free memory is underestimated in without compile case but is slightly overestimated in with compile case
4. We need to do two things: align the memory estimation between two cases and correctly estimate the memory in both cases
5. We can align the memory estimation in both cases by adding a call to empty_cache
6. We can make the memory estimation accurate by getting rid of the batch kv cache by setting self.kv_cache = [] before empty_cache call and by not adding batch_num_blocks to num_blocks.

[ajtejankar]

There's one more problem that the previous commit 03726e7 didn't fix. The mdoel_graph_wrapper object contains a reference to the kv_cache memory which prevents torch from freeing it up when empty_cache is called. Hence, we use model_graph_wrapper to estimate cuda graph memory overhead, delete it, and then re-instantiate one final time to actually warmup the graphs.

[tgaddair]

Nice, added a couple comments.

[tgaddair]

LGTM!