PegaFlow Python Package

High-performance key-value storage engine with Python bindings, built with Rust and PyO3.

Features

PegaEngine: Fast Rust-based key-value storage with Python bindings
PegaKVConnector: vLLM KV connector for distributed inference with KV cache transfer

Installation

From Source

# Install maturin if you haven't already
pip install maturin

# Build and install in development mode
cd python
maturin develop

# Or build a wheel
maturin build --release

From PyPI (coming soon)

pip install pegaflow

Usage

Basic KV Storage

from pegaflow import PegaEngine

# Create a new engine
engine = PegaEngine()

# Store key-value pairs
engine.put("name", "PegaFlow")
engine.put("version", "0.1.0")

# Retrieve values
name = engine.get("name")  # Returns "PegaFlow"
missing = engine.get("nonexistent")  # Returns None

# Remove keys
removed = engine.remove("name")  # Returns "PegaFlow"

vLLM KV Connector

from vllm import LLM
from vllm.distributed.kv_transfer.kv_transfer_agent import KVTransferConfig

# Configure vLLM to use PegaKVConnector
kv_transfer_config = KVTransferConfig(
    kv_connector="PegaKVConnector",
    kv_role="kv_both",
    kv_connector_module_path="pegaflow.connector",
)

# Create LLM with KV transfer enabled
llm = LLM(
    model="gpt2",
    kv_transfer_config=kv_transfer_config,
)

Connector Modes

PegaKVConnector defaults to read_write: it queries PegaFlow for reusable KV blocks, loads matched blocks into vLLM, and saves newly computed full blocks back to PegaFlow.

Set pegaflow.mode to save_only when another vLLM connector is responsible for reads and PegaFlow should only persist KV blocks for later reuse. This is intended for MultiConnector decode-side setups where an upstream connector owns the external hit/load path, while PegaFlow records the resulting KV cache. In save_only mode, PegaFlow does not query or load KV blocks.

vllm serve Qwen/Qwen3-0.6B \
  --kv-transfer-config '{
    "kv_connector": "MultiConnector",
    "kv_role": "kv_both",
    "kv_connector_extra_config": {
      "connectors": [
        {
          "kv_connector": "<external-read-connector>",
          "kv_role": "kv_both"
        },
        {
          "kv_connector": "PegaKVConnector",
          "kv_role": "kv_both",
          "kv_connector_module_path": "pegaflow.connector",
          "kv_connector_extra_config": {
            "pegaflow.mode": "save_only"
          }
        }
      ]
    }
  }'

Valid values are read_write and save_only.

TP Shards Across Hosts

CUDA IPC is host-local. When one tensor-parallel replica spans multiple hosts, run one PegaFlow server on each host and configure the connector with every server endpoint in global TP-rank order:

{
  "kv_connector": "PegaKVConnector",
  "kv_role": "kv_both",
  "kv_connector_module_path": "pegaflow.connector",
  "kv_connector_extra_config": {
    "pegaflow.tp_shard_endpoints": [
      "http://host-a:50055",
      "http://host-b:50055"
    ]
  }
}

For TP8 and two endpoints, global ranks 0-3 register with the first server and ranks 4-7 register with the second. Each server sees a local TP4 topology and must manage the four GPUs on its own host. Every vLLM process must receive the same ordered endpoint list.

The scheduler queries every shard and only reuses the prefix available from all of them. Each worker loads with the lease issued by its local server. The connector gives every shard a distinct namespace, so deployments with a different host split cannot reuse an incompatible cache layout.

TP sharding currently requires equal contiguous shards and TP-only parallelism. Pipeline, decode-context, and prefill-context parallelism are rejected when more than one endpoint is configured.

P/D Partial Tail Blocks

vLLM normally exposes hashes only for complete KV blocks. In a P/D deployment, enable pegaflow.pd_tail_save on prefill and pegaflow.pd_tail_load on decode to reuse the final partial prompt block as well. Start both vLLM processes with the same explicit PYTHONHASHSEED and --prefix-caching-hash-algo xxhash_cbor.

Prefill: {"pegaflow.pd_tail_save": true}

Decode: {"pegaflow.pd_tail_load": true, "pegaflow.wait_for_full_prefix": true}

pegaflow.wait_for_full_prefix makes decode wait (up to 30s) until the full prompt prefix is fetchable from a remote node via MetaServer + RDMA. It only applies when prefill and decode run separate engines; it does not observe saves landing in a shared/local engine and has no effect when RDMA is not configured.

Development

See the examples directory for more usage examples.

Testing

Running Unit Tests

The test suite includes integration tests that verify the EngineRpcClient can correctly communicate with a running pegaflow-server instance.

Prerequisites

Build the Rust extension:
```
cd python
maturin develop --release
```

Build the server binary:

cd ..
cargo build --release --bin pegaflow-server

Ensure CUDA is available (tests require GPU):

python -c "import torch; assert torch.cuda.is_available()"

Running Tests

cd python

# Run all tests
pytest tests/ -v

# Run specific test file
pytest tests/test_engine_client.py -v

# Run with coverage
pytest tests/ --cov=pegaflow --cov-report=html

Test Structure

tests/conftest.py: Contains pytest fixtures for:
- pega_server: Automatically starts/stops pegaflow-server for integration tests
- engine_client: Creates an EngineRpcClient connected to the test server
- client_context: Provides a ClientContext representing a vLLM instance with GPU KV cache tensors
- registered_instance: Provides a registered instance ID for query tests
tests/test_engine_client.py: Integration tests for:
- Server connectivity
- Query operations with various inputs

Test Fixtures

The ClientContext class abstracts a vLLM instance and provides:

register_kv_caches(): Register GPU KV cache tensors with the server
query(block_hashes): Query available blocks
unregister_context(): Unregister context from server

Example test usage:

def test_query(client_context):
    """Test query operation."""
    result = client_context.query([])
    assert result is not None

License

MIT

Metadata

Release files for pegaflow-llm 0.24.5

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Built distributions (wheels)

Table of built distributions (wheels) for pegaflow-llm 0.24.5
File	Interpreter	ABI	Platform	Reset
pegaflow_llm-0.24.5-cp314-cp314-manylinux_2_34_x86_64.whl 8.0 MB Sep 17, 2026 Without provenance	CPython 3.14	CPython 3.14	Linux glibc 2.34+ x86-64	Details
pegaflow_llm-0.24.5-cp313-cp313-manylinux_2_34_x86_64.whl 8.0 MB Sep 17, 2026 Without provenance	CPython 3.13	CPython 3.13	Linux glibc 2.34+ x86-64	Details
pegaflow_llm-0.24.5-cp312-cp312-manylinux_2_34_x86_64.whl 8.0 MB Sep 17, 2026 Without provenance	CPython 3.12	CPython 3.12	Linux glibc 2.34+ x86-64	Details
pegaflow_llm-0.24.5-cp311-cp311-manylinux_2_34_x86_64.whl 8.0 MB Sep 17, 2026 Without provenance	CPython 3.11	CPython 3.11	Linux glibc 2.34+ x86-64	Details
pegaflow_llm-0.24.5-cp310-cp310-manylinux_2_34_x86_64.whl 8.0 MB Sep 17, 2026 Without provenance	CPython 3.10	CPython 3.10	Linux glibc 2.34+ x86-64	Details

Total release size: 40.1 MB

Release files / pegaflow_llm-0.24.5-cp314-cp314-manylinux_2_34_x86_64.whl

Download URL	pegaflow_llm-0.24.5-cp314-cp314-manylinux_2_34_x86_64.whl
Size	8.0 MB
Tags	CPython 3.14 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`a4660e5881329a1bbbac275e25dcb9290afefed56226b36ff739d49239498985`
BLAKE2b-256 checksum How to use checksums	`8e95db221a8d366a5b476436300a8380bcc05e2c05d4111739b238d4c4c0abf7`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Release files / pegaflow_llm-0.24.5-cp313-cp313-manylinux_2_34_x86_64.whl

Download URL	pegaflow_llm-0.24.5-cp313-cp313-manylinux_2_34_x86_64.whl
Size	8.0 MB
Tags	CPython 3.13 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`3d63c6d9680417bdb5d3db020d4e423e50109a5c629083c3e911e36997f8b611`
BLAKE2b-256 checksum How to use checksums	`9604ef6abaa7894ad5ff6347c872913a4a621d72bc24db45cd72929edf4d0fb2`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Release files / pegaflow_llm-0.24.5-cp312-cp312-manylinux_2_34_x86_64.whl

Download URL	pegaflow_llm-0.24.5-cp312-cp312-manylinux_2_34_x86_64.whl
Size	8.0 MB
Tags	CPython 3.12 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`74ece1a2c7126218c9a126f3da14ec2416fbe98d34ef2084e85fedb452b59da7`
BLAKE2b-256 checksum How to use checksums	`eacf711f5870a309884c1dbc366f44d4c02dbdfa463a3a90ce58987521ca6e44`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Release files / pegaflow_llm-0.24.5-cp311-cp311-manylinux_2_34_x86_64.whl

Download URL	pegaflow_llm-0.24.5-cp311-cp311-manylinux_2_34_x86_64.whl
Size	8.0 MB
Tags	CPython 3.11 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`991cb5791d6fcfab2bbf749523f1d1346077276e429745fa4badc300c36db044`
BLAKE2b-256 checksum How to use checksums	`80d9b700acbdfd5b51f425e2f33505da9153a08e4689bd34cabecce4060b47af`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Release files / pegaflow_llm-0.24.5-cp310-cp310-manylinux_2_34_x86_64.whl

Download URL	pegaflow_llm-0.24.5-cp310-cp310-manylinux_2_34_x86_64.whl
Size	8.0 MB
Tags	CPython 3.10 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`5af311c5488296f0a8b8aa9d1f1976c800de7df15207a28db52d21050e02985c`
BLAKE2b-256 checksum How to use checksums	`6b549e7895664e220619616bdaf2dd1643391e2f7a05a37cb0eb08d91ab37932`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Security for pegaflow-llm 0.24.5

Report project as malware

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.