Running Large AI Models on Consumer Hardware

# Illustrative sketch (not from the paper)
import torch
from collections import OrderedDict

# Double‑buffer for two consecutive expert layers
class DoubleBuffer:
    def __init__(self, layer_fn):
        self.current = None          # GPU resident experts for current layer
        self.next = None             # Prefetched experts for next layer
        self.layer_fn = layer_fn

    def load_next(self, expert_ids):
        # Simulate async PCIe transfer (placeholder)
        self.next = torch.nn.ModuleList([self.layer_fn(eid) for eid in expert_ids])

    def swap(self):
        self.current, self.next = self.next, None

# Simple LRU cache for experts on GPU
class ExpertLRUCache:
    def __init__(self, capacity):
        self.capacity = capacity
        self.cache = OrderedDict()

    def get(self, eid):
        if eid in self.cache:
            self.cache.move_to_end(eid)
            return self.cache[eid]
        return None

    def put(self, eid, expert):
        self.cache[eid] = expert
        self.cache.move_to_end(eid)
        if len(self.cache) > self.capacity:
            self.cache.popitem(last=False)  # evict LRU

# Bandwidth‑adaptive execution (q* policy)
def execute_token(token, cache, bandwidth_pcie, bandwidth_cpu):
    # Decide where to fetch missing expert
    if bandwidth_pcie >= bandwidth_cpu:
        # Prefer PCIe fill then GPU exec
        expert = cache.get(token.expert_id) or load_expert_via_pcie(token.expert_id)
    else:
        # Direct CPU execution fallback
        expert = load_expert_via_cpu(token.expert_id)
    # Run expert (placeholder)
    return expert(token.input)

# Placeholder loaders
def load_expert_via_pcie(eid):
    # In real system this streams over PCIe into GPU memory
    return torch.nn.Linear(1024, 1024)

def load_expert_via_cpu(eid):
    # CPU‑only execution path
    return lambda x: x  # no‑op for illustration

# Example workflow for a single decode step
def decode_step(tokens, double_buf, lru_cache, bw_pcie, bw_cpu):
    # Prefetch next layer experts while current layer runs
    double_buf.load_next([t.expert_id for t in tokens])
    outputs = []
    for token in tokens:
        expert = lru_cache.get(token.expert_id)
        if expert is None:
            expert = execute_token(token, lru_cache, bw_pcie, bw_cpu)
            lru_cache.put(token.expert_id, expert)
        outputs.append(expert(token.input))
    double_buf.swap()  # hide latency for next layer
    return outputs

// Illustrative sketch (not from the paper)
const { EventEmitter } = require('events');

// Double‑buffer for two consecutive expert layers
class DoubleBuffer extends EventEmitter {
  constructor(layerFactory) {
    super();
    this.current = null; // GPU resident experts for current layer
    this.next = null;    // Prefetched experts for next layer
    this.layerFactory = layerFactory;
  }

  loadNext(expertIds) {
    // Simulate async PCIe transfer (placeholder)
    this.next = expertIds.map(id => this.layerFactory(id));
    this.emit('prefetchDone');
  }

  swap() {
    this.current = this.next;
    this.next = null;
  }
}

// Simple LRU cache for experts on GPU
class ExpertLRUCache {
  constructor(capacity) {
    this.capacity = capacity;
    this.map = new Map(); // preserves insertion order
  }

  get(eid) {
    if (!this.map.has(eid)) return null;
    const val = this.map.get(eid);
    this.map.delete(eid);
    this.map.set(eid, val); // move to recent
    return val;
  }

  put(eid, expert) {
    if (this.map.has(eid)) this.map.delete(eid);
    this.map.set(eid, expert);
    if (this.map.size > this.capacity) {
      // evict LRU (first entry)
      const lruKey = this.map.keys().next().value;
      this.map.delete(lruKey);
    }
  }
}

// Bandwidth‑adaptive execution (q* policy)
function executeToken(token, cache, bwPcie, bwCpu) {
  let expert;
  if (bwPcie >= bwCpu) {
    expert = cache.get(token.expertId) || loadExpertViaPCIe(token.expertId);
  } else {
    expert = loadExpertViaCPU(token.expertId);
  }
  // Run expert (placeholder)
  return expert(token.input);
}

function loadExpertViaPCIe(eid) {
  // In real system this streams over PCIe into GPU memory
  return input => input; // no‑op placeholder
}

function loadExpertViaCPU(eid) {
  // CPU‑only execution path
  return input => input; // no‑op placeholder
}

// Example workflow for a single decode step
function decodeStep(tokens, doubleBuf, lruCache, bwPcie, bwCpu) {
  doubleBuf.loadNext(tokens.map(t => t.expertId));
  const outputs = [];
  for (const token of tokens) {
    let expert = lruCache.get(token.expertId);
    if (!expert) {
      expert = executeToken(token, lruCache, bwPcie, bwCpu);
      lruCache.put(token.expertId, expert);
    }
    outputs.push(expert(token.input));
  }
  doubleBuf.swap(); // hide latency for next layer
  return outputs;
}