[Android] initial UpdateLogitsOrProbOnCPUSync freezes OS

[tobrun]

🐛 Bug

👋 been diving more into the freezing issue noted in #1379, creating a new issue with better details and ways to reproduce.

When I load a llama-2-7b model (any of the default configured quantizations) and prompt the model with a large input:

• Android OS freezes and becomes unresponsive
• System application crashes and restarts itself
• The output of the LLM starts appearing in the MLC-Chat app after the os becomes responsive again.

Additional information: prompt size ~ 2700 characters (which is gets truncated: The prompt tokens are more than max_window_size, the input will be truncated.) which is not the root cause but makes it easier to reproduce.

---

When timing execution of all the different functions in llm_chat.cc, I can see that the underlying issue is coming from:

SampleTokenFromLogits executed in 65189.915079 ms

Diving more deep into where it's being hold back, it's coming from:

  void UpdateLogitsOrProbOnCPUSync(NDArray logits_or_prob) {    
    if (!logits_on_cpu_.defined()) {
      logits_on_cpu_ = logits_or_prob.CopyTo(DLDevice{kDLCPU, 0});
    } else {
      ICHECK_EQ(logits_on_cpu_->shape[0], logits_or_prob->shape[0])
          << "Expect size of logits remain unchanged";
      logits_on_cpu_.CopyFrom(logits_or_prob);
    }
    TVMSynchronize(device_.device_type, device_.device_id, nullptr);
  }

The next line when logits aren't found on the CPU is the last line to hit before it halts:

logits_on_cpu_ = logits_or_prob.CopyTo(DLDevice{kDLCPU, 0});

Will continue to look more into what above line does but any 👀 are highly appreciated!

[tobrun]

Smaller model like https://github.com/jzhang38/TinyLlama, is faster but still results in a freeze of the OS:

SampleTokenFromLogits executed in 12231.236818 ms

using a single world prompt, like hello, results in:

SampleTokenFromLogits executed in 567.032344 ms

While it's shorter, having any freezes at all makes it unusable in non-chat UI use-cases. For example, using a with voice assistant powered by a LLM while having animating elements and sound in the application.

[tqchen]

Would be great to confirm if the particular run executes on a separate thread.

The main issue is not sample, but the cost of Synchronize(aka the cost of waiting for GPU compute to finish)

[tobrun]

Would be great to confirm if the particular run executes on a separate thread.

it is, my initial assumption about main thread halting was wrong.
The code is executing on a separate thread.

[tqchen]

I see, then it might due to the fact that we are overusing the gpu resources for prefill.

[tqchen]

We are moving towards a new engine and android sdk, so this issues as of now is less relevant, closing for now. We might still face some throttling issue in new engine depending on amount of compute

[grabbou]

Unfortunately, I am still observing this behavior on Galaxy S23, and v19 branch. First time running model freezes everything for ~5/6 seconds, subsequent runs are sub second, but still. I would be happy to provide necessary feedback.

[artus9033]

+1, I am experiencing this as well. It seems that the whole OS is freezing, so most likely this is some excessive usage of resources / heavy allocation going on. Would you be able to look into this @tqchen ?

[artus9033]

For better visibility, I raised a new issue, #3363, as this one seems to be inactive with no perspective for a re-open.