SSD-Pooled Memory Serving for Large-Scale AI Workloads
MemServe is a GPU–DRAM–SSD hierarchical memory-serving prototype for AI workloads whose working sets exceed GPU memory capacity.
Currently supported workloads:
- GNN feature serving
- MoE expert serving
GNN MoE
│ │
Feature Request Expert Request
│ │
└─────────┬─────────┘
▼
┌──────────┐
│ MemServe │
└────┬─────┘
│
┌────────────┼────────────┐
▼ ▼ ▼
GPU HBM Host DRAM SSD Pool
Tier 0 Tier 1 Tier 2
Full design writeup: docs/architecture.md.
- Unified
MemoryObjectabstraction shared by every workload - GPU → DRAM → SSD hierarchical serving with promote/demote on eviction
- Workload-aware caching: eviction priority is computed by workload- specific code, consumed as an opaque score by workload-agnostic mechanism (see Design below)
- Parallel SSD pooling with single / round-robin / load-aware placement
- Trace-driven evaluation with a queueing-based SSD timing model
All numbers below are from results/gnn.csv and results/moe.csv,
regenerated by benchmarks/run_all.py. Full plots in
results/figures/.
GPU memory pressure is where workload-aware caching earns its keep. At 10% GPU capacity, workload-aware caching lowers estimated MoE runtime by ~3% vs LRU; at 40% GPU capacity the gap widens to ~22% (17.8s vs 22.6s of I/O stall). GNN shows the same direction of effect but a smaller magnitude, consistent with its access pattern being less skewed than MoE's expert popularity distribution.
SSD pooling materially reduces I/O stall under contention. Moving from 1 SSD to a 2-device load-aware pool cuts MoE I/O stall from 86.9s to 31.9s (~2.7x) on the same trace; GNN sees stall drop from 3.75s to 1.30s (~2.9x). A 4-device pool gives further but diminishing returns.
Workload-aware caching reduces SSD traffic, not just latency. For MoE, workload-aware caching cuts SSD bytes transferred from ~126GB (LRU) to ~92GB on the same trace and capacity — the policy is avoiding re-fetches, not just serving them faster.
| Figure | Question it answers |
|---|---|
fig1_gpu_capacity.png |
Does workload-aware caching help more when GPU memory is tight? |
fig2_ssd_pooling.png |
Does SSD pooling actually reduce I/O stall? |
fig3_policy_comparison.png |
Why do we need workload-aware memory management at all? |
GNN mini-batch training and MoE inference both regularly need to touch more data per step than fits in GPU memory — node features for neighbor-sampled subgraphs, or expert weight shards selected by a gating network. Both problems reduce to the same question: given a capacity-bounded GPU/DRAM cache and a slower SSD backing store, which objects do you keep resident, and how do you get to the SSD store fast when you have to? MemServe treats that as one problem instead of two separate, workload-specific caching layers.
The core principle: workload semantics are separated from memory- management mechanisms.
memserve/ never branches on whether an object is a GNN feature or an
MoE expert. It only sees a MemoryObject with a priority field, and
its eviction policies (memserve/cache.py) just evict
min(priority). Each workload (workloads/gnn, workloads/moe)
supplies its own priority function, which is the only place workload-
specific intelligence enters the system.
# memserve/cache.py -- mechanism, workload-agnostic
class WorkloadAwarePolicy(EvictionPolicy):
def choose_victim(self, resident):
return min(resident.values(), key=lambda o: o.priority).id
# workloads/gnn/policy.py -- semantics, workload-specific
def gnn_priority(obj: MemoryObject) -> float:
return float(obj.access_count) + 0.001 * obj.last_accessAdding a third workload means writing a new workloads/<name>/ module.
It requires zero changes to memserve/.
The public API surface workload code touches is four calls:
request(), prefetch(), promote(), evict() — see
docs/architecture.md.
benchmarks/run_all.py runs every experiment behind the three figures
and writes results/*.csv + results/figures/*.png. See
docs/architecture.md
for how the SSD timing model works — in particular, why it uses
Poisson-distributed arrivals rather than a fixed interval.
git clone <this-repo>
cd MemServe
pip install -r requirements.txt
python examples/run_gnn.py
python examples/run_moe.pyExpected output (run_gnn.py):
Workload: GNN (feature serving)
Requests: 100000
GPU hits: 73262 (73.26%)
DRAM hits: 8938 (8.94%)
SSD reads: 17800 (17.80%)
Transferred: 0.01 GB
I/O stall: 1.78 s
Estimated runtime: 2.11 s
python benchmarks/run_all.pyThis regenerates results/gnn.csv, results/moe.csv, and all three
figures in results/figures/ from scratch — no external datasets
required, since traces are generated synthetically (see
docs/limitations.md).
Run the test suite:
python -m pytest tests/ -vMemServe's feature-caching problem is related to prior work on SSD-backed GNN training systems (e.g. Ginex-style feature caching) and to MoE serving systems that manage expert placement across memory tiers. MemServe's contribution relative to that line of work is the shared abstraction across both workload families, rather than a new per-workload caching algorithm — see docs/architecture.md.
MemServe is a trace-driven prototype, not a production runtime. See docs/limitations.md for an explicit list of what's really implemented vs. simulated.
See docs/project-evolution.md for how this generalized from an earlier SSD-backed GNN investigation.
MemServe/
├── memserve/ # core library: workload-agnostic mechanism
├── workloads/ # GNN and MoE: workload-specific semantics
├── benchmarks/ # experiment driver -> results/*.csv, figures
├── examples/ # 5-minute demo entry points
├── tests/ # unit tests (pytest)
├── docs/ # architecture, limitations, project history
└── results/ # generated CSVs and figures