flexFlow — Adaptive KV Cache Transfer for Distributed LLM Inference

flexFlow is a lightweight middleware that dynamically selects the optimal data transfer path for KV cache migration in distributed LLM serving systems (e.g., vLLM, SGLang with disaggregated prefill/decode). Instead of always using the same direct link, flexFlow monitors real-time link utilization and routes transfers through whichever path delivers the best throughput — direct NVLink, relay via an idle GPU, cross-node RDMA, or even shared storage — all while respecting QoS constraints to avoid interfering with latency-critical inference traffic.

Why flexFlow?

In production LLM serving clusters, KV cache transfers compete with model communication (tensor parallelism AllReduce, pipeline parallelism P2P) for the same NVLink and RDMA bandwidth. When a direct link is congested:

• Naive approach: transfers queue behind inference traffic, increasing Time-To-First-Token (TTFT)
• flexFlow approach: detects congestion in ~15 microseconds, reroutes through idle links, and maintains high throughput

Measured results on DGX A100 (4 GPUs, NVLink):

Scenario	Always-Direct	flexFlow Adaptive	Improvement
Idle system	22.53 GB/s	22.53 GB/s	1.00x
NVLink contention (85%)	11.68 GB/s	21.99 GB/s	1.88x

Cross-node validation on DGX V100 (2 nodes, InfiniBand):

Metric	Result
Cross-node RDMA throughput	0.38 GB/s (UCX/TCP)
Data integrity verification	PASS
PathSelector decision latency	14.6 us
Congestion-aware path switching	CrossNodeRelay -> StorageRelay

Key Features

• Sub-100us decision latency — PathSelector runs entirely on CPU, no GPU sync required
• 5 path types — Direct, NVLinkRelay, CrossNodeRelay (RDMA), SplitPath (parallel multi-path), StorageRelay
• QoS-aware — configurable utilization thresholds prevent interference with inference traffic on NVLink, RDMA, CNIC, and SNIC links
• TTFT-sensitive mode — penalizes multi-hop paths when first-token latency is critical
• SplitPath — automatically splits data across two link-disjoint paths when congestion makes it beneficial
• Hardware auto-discovery — detects GPU topology via NVML, RDMA NICs via sysfs
• Cross-node RDMA — GPU-to-GPU transfers via NIXL (NVIDIA's transfer library) over InfiniBand/RoCE
• Mock mode — full decision logic works without GPUs for development and testing

Quick Start

Build

# Prerequisites: CMake >= 3.18, C++17 compiler, Python >= 3.8, pybind11
pip install pybind11 torch

# Build C++ core + Python extension
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(nproc)
cd ..

# Install Python package
pip install -e .

Run the mock demo (no GPU required)

python examples/adaptive_transfer_demo.py --mock

This walks through every API step: topology creation, monitoring setup, path selection under idle/contention/TTFT-sensitive scenarios, and decision latency benchmarking.

Run with real GPUs (>= 3 GPUs with NVLink)

python examples/adaptive_transfer_demo.py --size_mb 64 --iters 50

Cross-node demo (2 nodes with RDMA)

srun -N 2 --ntasks=2 --ntasks-per-node=1 --gres=gpu:8 -t 0:15:00 \
    bash scripts/run_cross_node_demo.sh

Usage

Basic API

from flexflow import Topology, NodeAgent, PathSelector, TransferRequest, SelectorConfig

# 1. Build topology (auto-discover or mock)
topo = Topology()
topo.discover()           # Real hardware via NVML + sysfs
# topo = Topology.make_mock_dgx(8)  # Or mock for testing

# 2. Start monitoring
agent = NodeAgent(topo, sample_interval_ms=100)
agent.start()  # Background thread polls NVLink/PCIe/RDMA utilization

# 3. Configure and create path selector
cfg = SelectorConfig()
cfg.relay_benefit_margin = 0.10   # Relay must beat direct by 10%
cfg.enable_split_path = True      # Allow parallel multi-path transfers
selector = PathSelector(topo, agent, cfg)

# 4. Select best path for a transfer
gpus = topo.gpu_nodes()
req = TransferRequest(src=gpus[0], dst=gpus[7], data_bytes=64*1024*1024)
decision = selector.select(req)

print(decision.chosen_type)        # Direct, NVLinkRelay, SplitPath, etc.
print(decision.estimated_bandwidth) # GB/s
print(decision.hops)                # Node sequence: [src, ..., dst]

# 5. Execute the transfer
from flexflow.nixl_adapter import TorchAdapter
adapter = TorchAdapter()
result = adapter.execute(decision, src_tensor, dst_tensor, relay_tensor)
print(f"{result.throughput_gbps:.2f} GB/s")

Cross-Node RDMA Transfer

from flexflow.nixl_adapter import CrossNodeNixlAdapter, get_nodelist

# Both ranks run this
adapter = CrossNodeNixlAdapter(rank=rank, nodes=get_nodelist(), gpu_device=0)
adapter.setup(tensor)  # Exchanges NIXL metadata via TCP

# Rank 1 reads from Rank 0
result = adapter.rdma_transfer(tensor)
print(f"{result.throughput_gbps:.2f} GB/s")

adapter.close()

Tunable Configuration

cfg = SelectorConfig()
# Scoring weights
cfg.w_bandwidth = 1.0          # Throughput benefit weight
cfg.w_queue_delay = 0.5        # Congestion penalty weight
cfg.w_hops = 0.2               # Hop count penalty weight
cfg.w_qos_risk = 0.3           # QoS interference risk weight
cfg.w_latency = 0.5            # TTFT latency penalty weight

# QoS thresholds (reject paths with links above these)
cfg.nvlink_util_threshold = 0.80
cfg.rdma_util_threshold = 0.80
cfg.cnic_util_threshold = 0.70
cfg.snic_util_threshold = 0.85

# SplitPath
cfg.enable_split_path = True
cfg.max_split_paths = 2
cfg.split_benefit_margin = 0.15  # Split must beat single by 15%

Applicable Scenarios

Scenario	How flexFlow Helps
Disaggregated prefill/decode	KV cache migrated between prefill and decode GPU groups; flexFlow avoids congested NVLink paths
Multi-node KV cache sharing	Cross-node RDMA transfers routed around congested NICs
Prompt caching across nodes	Shared storage relay when RDMA is saturated
Mixed workloads	QoS thresholds protect inference traffic from bulk KV transfers
TTFT-critical requests	Low-hop path preference minimizes first-token latency overhead

Performance Data

Metric	Value	Hardware
PathSelector.select() latency	14-15 us	Any CPU
PathSelector.select() P99	< 50 us	Any CPU
Throughput (idle NVLink)	22.5 GB/s	A100 SXM4
Throughput (contention, adaptive)	22.0 GB/s	A100 SXM4
Throughput (contention, direct)	11.7 GB/s	A100 SXM4
Cross-node RDMA	0.38 GB/s	V100 DGX (UCX/TCP)
Decision overhead vs transfer time	< 0.1%	64 MB transfer

Project Layout

src/
  core/
    types.h/.cc          Enums, structs (NodeType, LinkType, PathType, PathDecision, ...)
    topology.h/.cc       Hardware graph: nodes, links, path enumeration, mock factories
    node_agent.h/.cc     Background monitoring thread: NVLink, PCIe, RDMA utilization
    path_selector.h/.cc  Scoring engine: select(), score(), QoS constraints, SplitPath
  monitor/
    nvml_monitor.h/.cc   NVML wrapper: GPU info, NVLink peers/utilization, PCIe throughput
    sysfs_monitor.h/.cc  Sysfs wrapper: RDMA NIC discovery and byte counters
  python/
    bindings.cc          pybind11 bindings for all C++ types and classes

python/flexflow/
    __init__.py          Package exports
    nixl_adapter.py      Transfer adapters: NixlAdapter, TorchAdapter, CrossNodeNixlAdapter
    benchmark.py         Benchmark utilities (BenchmarkResult, measure_selector_latency)

tests/
    test_topology.cc     C++ topology tests
    test_path_selector.cc C++ path selector tests
    python/              Python test suite (pytest)

benchmarks/              Performance benchmarks (selector latency, throughput, contention)
examples/                End-to-end demos (single-node, cross-node, adaptive)

Testing

# C++ tests
cd build && ctest --output-on-failure

# Python tests
python -m pytest tests/python/ -v

# Mock demo (no GPU)
python examples/adaptive_transfer_demo.py --mock

License

Apache 2.0