[Question]: Out of memory for MInference in e2e benchmark test & NIAH with long context_window

Describe the issue

Hi, thank you for your excellent work on MInference and for open-sourcing such a well-structured and impactful repository.

I'm trying to reproduce the end-to-end benchmark experiments from this repository, but I'm encountering a CUDA out-of-memory error (with a 80G A100 GPU).

For e2e benchmark with target length the following command, the OOM error occurs

python experiments/benchmarks/benchmark_e2e.py --attn_type minference --context_window 1_000_000 --kv_cache_cpu

The log is:

CUDA_VISIBLE_DEVICES=0 python experiments/benchmarks/benchmark_e2e.py --attn_type minference --context_window 1_000_000 --kv_cache_cpu
/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/utils/hub.py:111: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.
  warnings.warn(
INFO 12-02 19:13:31 [__init__.py:243] Automatically detected platform cuda.
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████| 4/4 [00:02<00:00,  1.84it/s]
<---- MInference Config Detail ----> attn_type minference, kv_type dense
Traceback (most recent call last):
  File "/home/cpwu/working/SPTCAttn/MInference/experiments/benchmarks/benchmark_e2e.py", line 147, in <module>
    run_target_length(args.context_window, model, args.attn_type)
  File "/home/cpwu/working/SPTCAttn/MInference/experiments/benchmarks/benchmark_e2e.py", line 33, in run_target_length
    model(
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 969, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 688, in forward
    outputs: BaseModelOutputWithPast = self.model(
                                       ^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 969, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 453, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 48, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/minference/patch.py", line 527, in forward_llama_decoder_layer
    attention_outputs = self.self_attn(
                        ^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/minference/patch.py", line 874, in <lambda>
    lambda self, *args, **kwargs: forward(self, *args, **kwargs)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/minference/modules/forward.py", line 85, in attn_forward
    query_states, key_states = apply_rotary_pos_emb(
                               ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cpwu/working/SPTCAttn/MInference/.venv/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 145, in apply_rotary_pos_emb
    q_embed = (q * cos) + (rotate_half(q) * sin)
                           ~~~~~~~~~~~~~~~^~~~~
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 7.63 GiB. GPU 0 has a total capacity of 79.25 GiB of which 2.23 GiB is free. Including non-PyTorch memory, this process has 77.02 GiB memory in use. Of the allocated memory 65.06 GiB is allocated by PyTorch, and 11.46 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Similarly, for all latency experiments, the OOM also occurs for context_window of 500k and 1000k (success for 1k~100k), with command

python experiments/benchmarks/benchmark_e2e.py --run_benchmark

And in the NIAH test, OOM also occurs (successful for jobs 0-4, OOM for jobs 4-15)

Start with experiments/benchmarks/benchmark_e2e.py, especially run_target_length, and trace the failing call into minference/modules/forward.py at attn_forward. Reproduce the stated 1,000,000-token command and the latency or NIAH cases, then determine whether the memory behavior is expected or requires a change. Done means documenting a verified workaround or fixing the out-of-memory behavior across the reported context windows.