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}
Partial-tail caching requires a dense attention layout and is rejected for HMA or SlidingWindow cache groups.
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/stopspegaflow-serverfor integration testsengine_client: Creates anEngineRpcClientconnected to the test serverclient_context: Provides aClientContextrepresenting a vLLM instance with GPU KV cache tensorsregistered_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 serverquery(block_hashes): Query available blocksunregister_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
1 maintainer
Metadata
1 maintainer
Release files for pegaflow-llm 0.25.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distributions (wheels)
| File | Reset | |||
|---|---|---|---|---|
| pegaflow_llm-0.25.1-cp314-cp314-manylinux_2_34_x86_64.whl | CPython 3.14 | CPython 3.14 | Linux glibc 2.34+ x86-64 | Details |
| pegaflow_llm-0.25.1-cp313-cp313-manylinux_2_34_x86_64.whl | CPython 3.13 | CPython 3.13 | Linux glibc 2.34+ x86-64 | Details |
| pegaflow_llm-0.25.1-cp312-cp312-manylinux_2_34_x86_64.whl | CPython 3.12 | CPython 3.12 | Linux glibc 2.34+ x86-64 | Details |
| pegaflow_llm-0.25.1-cp311-cp311-manylinux_2_34_x86_64.whl | CPython 3.11 | CPython 3.11 | Linux glibc 2.34+ x86-64 | Details |
| pegaflow_llm-0.25.1-cp310-cp310-manylinux_2_34_x86_64.whl | CPython 3.10 | CPython 3.10 | Linux glibc 2.34+ x86-64 | Details |
Total release size: 40.9 MB
Release files / pegaflow_llm-0.25.1-cp314-cp314-manylinux_2_34_x86_64.whl
| Download URL | pegaflow_llm-0.25.1-cp314-cp314-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 8.2 MB |
| Tags | CPython 3.14 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
2c2a1f21f41e50b0ae66ac83322c84c0ad86d3d893e9c658e5173eac476ffd7d
|
|
BLAKE2b-256 checksum How to use checksums |
1a88e50e7f964ba3a74d16cb10d145141a96a79e7491345f86dd1a8ba3d21963
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Release files / pegaflow_llm-0.25.1-cp313-cp313-manylinux_2_34_x86_64.whl
| Download URL | pegaflow_llm-0.25.1-cp313-cp313-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 8.2 MB |
| Tags | CPython 3.13 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
e7dfc9df2bfa8ccf918f536dbe2f2007dc9dc5b5a2d89fc63bcfd9dc84804e3a
|
|
BLAKE2b-256 checksum How to use checksums |
c87b0eb5eb9a5b30bc8404544411bcf65a3e930dac12d74979c73cec27eed149
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Release files / pegaflow_llm-0.25.1-cp312-cp312-manylinux_2_34_x86_64.whl
| Download URL | pegaflow_llm-0.25.1-cp312-cp312-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 8.2 MB |
| Tags | CPython 3.12 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
63377579d217ad0289609f2abafdd5304d88c875862ef8b9f261c88dc863efec
|
|
BLAKE2b-256 checksum How to use checksums |
af580487f4026226c9373a1d7de47a7f5b7bcb305f2a003b8e6d8e9841409694
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Release files / pegaflow_llm-0.25.1-cp311-cp311-manylinux_2_34_x86_64.whl
| Download URL | pegaflow_llm-0.25.1-cp311-cp311-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 8.2 MB |
| Tags | CPython 3.11 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
4ac2e62bc4fe4f42f41a3e690dbd0c8b38fd8ecc4ffc91b710cebe3d0563372c
|
|
BLAKE2b-256 checksum How to use checksums |
a4afceb316d739110711b46033df08750e15bb47efa9d204389a1567ad2d578c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Release files / pegaflow_llm-0.25.1-cp310-cp310-manylinux_2_34_x86_64.whl
| Download URL | pegaflow_llm-0.25.1-cp310-cp310-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 8.2 MB |
| Tags | CPython 3.10 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
326eb19533511052e928e75064e9980fdb3f7cc315bb49e7013faee7f365f82f
|
|
BLAKE2b-256 checksum How to use checksums |
047c8e0acec2381564e12181485c2454a39b53b7529ff2b65484fc9f38ffe288
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Security for pegaflow-llm 0.25.1
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.
