# prain

Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places
each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk.
Missing experts are copied to the GPU or computed on the CPU (`[miss]`). Eviction, heat, prefetch
and miss policies are configurable, and `prain tune` picks them from a routing trace.
Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 GB RX 6600 XT.

## Build (GPU host, user space)

    scripts/build.sh wrappers sdk fetch torch torchvision vllm prain

- `sdk`: venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
- `torch`, `torchvision`: from source with clang and `-march=native` (hours)
- `vllm`: v0.32.0 from source with gfx1032 added and `patches/vllm-*.patch` applied
- `prain`: editable install and the test suite

Check the stack: `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`

## Serve

    . scripts/env.sh
    prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --served-model-name gemma4 \
        --language-model-only --max-model-len 4096 --max-num-batched-tokens 1024 \
        --enable-prefix-caching --attention-backend TRITON_ATTN

The first start writes the pack (~3 min). Chat with `prain tui`, or `prain web --config FILE` for
several users with drafts.

## Tune for a machine

    python scripts/measure.py PACK_DIR/layer029.w13 --threads 16 > cost.json   # server stopped
    prain sim TRACE CONFIG MODEL_DIR/config.json --cost cost.json
    prain tune TRACE MODEL_DIR/config.json --cost cost.json --expert-vram-gib 1.67 --ram-gib 0 \
        --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
    python scripts/bench.py http://127.0.0.1:8000

The simulator treats every disk-tier expert as a disk read; the page cache makes real runs faster.

## vLLM patches

- `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"
  checks took it for Blackwell (B200 attention tiles: 128 KiB of LDS on a 64 KiB GPU).
- `vllm-0.32-gemma4-fit-8gb`: Gemma 4 shares `embed_tokens` as `lm_head` instead of allocating a
  second 1.4 GiB copy while loading, and sizes rope caches by `max_model_len` (saves ~0.4 GiB).