core/nn/llama_streaming library

AirLLM-style layer-streaming Llama runner.

Instead of holding all L transformer layers in memory at once, this runner keeps:

  • the token embedding (embed_tokens.weight),
  • the final RMSNorm,
  • the untied lm_head (if any),
  • a single resident LlamaBlock,
  • a RopeCache.

During forward, each layer's parameters are pulled from disk via a ShardedSafeTensorsReader, adoptCpuStorageFrom'd into the resident block's tensors, and the block is run once. Peak resident memory ≈ embed table + one layer's weights + activations.

This is the trick AirLLM uses to fit 70B models on a 4 GB card: per-token latency is dominated by disk→RAM bandwidth × numLayers, but the model fits.

Constraints:

  • CPU only. GPU streaming would require a re-upload per layer, which the underlying Tensor.adoptCpuStorageFrom fast path does not support.
  • fp16 checkpoints strongly preferred — the fp16 fast path in the loader avoids fp32 promotion, halving both disk I/O and resident layer bytes.
  • No KV cache in this first cut. Every forward re-processes the full prefix. For research / porting demos this is fine; for production autoregression a resident-per-layer KV cache is the natural next step.

final reader = ShardedSafeTensorsReader.open('model.safetensors'); final runner = LlamaStreamingRunner(cfg, reader); final out = runner.generate(promptIds, maxNewTokens: 20);

Functions

estimateLayerBytes(ShardedSafeTensorsReader reader, LlamaConfig cfg, {int layer = 0}) → int
Sum of on-disk weight bytes for one layer, computed from the reader's header (no tensor decoding). Useful for planning: multiply by the number of resident layers you want to fit.