On-Device LLM Inference at 50M Mobile Sessions: llama.cpp vs MLX vs MediaPipe vs ExecuTorch

#include "llama.h"
#include <chrono>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

struct ModelDeleter { void operator()(llama_model* m) const { llama_model_free(m); } };
struct CtxDeleter { void operator()(llama_context* c) const { llama_free(c); } };
using ModelPtr = std::unique_ptr<llama_model, ModelDeleter>;
using CtxPtr = std::unique_ptr<llama_context, CtxDeleter>;

struct BackendGuard {
  BackendGuard() { llama_backend_init(); }
  ~BackendGuard() { llama_backend_free(); }
};

struct DecodeStats { double prefill_ms; double decode_tok_s; int generated; };

DecodeStats run_prefill_decode(const std::string& model_path, const std::string& prompt,
                               int n_predict, int n_threads) {
  BackendGuard backend;
  llama_model_params mp = llama_model_default_params();
  mp.n_gpu_layers = 99; // Metal/Vulkan where available; fallback to CPU if unsupported
  ModelPtr model(llama_model_load_from_file(model_path.c_str(), mp));
  if (!model) throw std::runtime_error("model load failed");

  llama_context_params cp = llama_context_default_params();
  cp.n_ctx = 2048;
  cp.n_batch = 512;
  cp.n_threads = n_threads;
  CtxPtr ctx(llama_init_from_model(model.get(), cp));
  if (!ctx) throw std::runtime_error("context init failed");

  const llama_vocab* vocab = llama_model_get_vocab(model.get());
  int n = -llama_tokenize(vocab, prompt.c_str(), (int32_t)prompt.size(), nullptr, 0, true, true);
  if (n <= 0) throw std::runtime_error("tokenize sizing failed");
  std::vector<llama_token> tokens(n);
  if (llama_tokenize(vocab, prompt.c_str(), (int32_t)prompt.size(), tokens.data(), n, true, true) < 0)
    throw std::runtime_error("tokenize failed");

  auto t0 = std::chrono::steady_clock::now();
  llama_batch batch = llama_batch_get_one(tokens.data(), n);
  if (llama_decode(ctx.get(), batch) != 0) throw std::runtime_error("prefill failed");
  auto t1 = std::chrono::steady_clock::now();

  int generated = 0;
  for (; generated < n_predict; ++generated) {
    const float* logits = llama_get_logits_ith(ctx.get(), -1);
    llama_token next = 0;
    for (int32_t t = 1; t < llama_vocab_n_tokens(vocab); ++t)
      if (logits[t] > logits[next]) next = t;
    if (llama_vocab_is_eog(vocab, next)) break;
    llama_batch one = llama_batch_get_one(&next, 1);
    if (llama_decode(ctx.get(), one) != 0) break;
  }
  auto t2 = std::chrono::steady_clock::now();
  double prefill_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
  double decode_s = std::chrono::duration<double>(t2 - t1).count();
  return {prefill_ms, decode_s > 0 ? generated / decode_s : 0.0, generated};
}

import com.google.mediapipe.tasks.genai.llminference.LlmInference
import java.nio.file.Files
import java.nio.file.Paths

fun createLlm(context: android.content.Context, modelPath: String): LlmInference {
    require(Files.exists(Paths.get(modelPath))) { "Missing model: $modelPath" }
    val options = LlmInference.LlmInferenceOptions.builder()
        .setModelPath(modelPath)
        .setMaxTokens(512)
        .setTopK(40)
        .setTemperature(0.8f)
        .build()
    return LlmInference.createFromOptions(context, options)
}

import torch
from executorch.exir import to_edge_transform_and_lower
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner

class CausalLM(torch.nn.Module):
    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model(input_ids).logits

try:
    ep = torch.export.export(CausalLM(model).eval(), (torch.ones(1, 128, dtype=torch.long),))
    edge = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()])
    edge.to_executorch().write_to_file("gemma-2b.pte")
except Exception as exc:
    raise RuntimeError("ExecuTorch export failed") from exc

fun installThermalGovernor(context: android.content.Context, session: SessionController) {
    val pm = context.getSystemService(android.os.PowerManager::class.java)
    pm.addThermalStatusListener(context.mainExecutor) { status ->
        session.maxNewTokens = when (status) {
            android.os.PowerManager.THERMAL_STATUS_NONE,
            android.os.PowerManager.THERMAL_STATUS_LIGHT -> 512
            android.os.PowerManager.THERMAL_STATUS_MODERATE -> 256
            android.os.PowerManager.THERMAL_STATUS_SEVERE -> 128
            else -> 0 // route to cloud or stop generation cleanly
        }
        session.telemetry.record("thermal_status", status)
    }
}

final class ThermalGovernor {
    private var observer: NSObjectProtocol?
    private(set) var state: ProcessInfo.ThermalState = .nominal

    func start(onChange: @escaping (ProcessInfo.ThermalState) -> Void) {
        observer = NotificationCenter.default.addObserver(
            forName: ProcessInfo.thermalStateDidChangeNotification,
            object: nil,
            queue: .main
        ) { [weak self] _ in
            guard let self else { return }
            self.state = ProcessInfo.processInfo.thermalState
            onChange(self.state)
        }
    }

    deinit {
        if let observer { NotificationCenter.default.removeObserver(observer) }
    }
}