Skip to content

NeuralOps v1.3.4: Serving Primitives & KV Cache Optimization

Choose a tag to compare

View all tags
@tahamajs tahamajs released this 17 Oct 22:07
· 7803 commits to main since this release
v1.3.4
a70ea50

🚀 NeuralOps v1.3.4: Serving Primitives & KV Cache Optimization

Non-Contiguous Paged KV Cache & Speculative Decoding Prototypes

PagedKVCache · Zero Memory Fragmentation · Speculative Decoding Engine · Continuous Batching Primitives

PyPI version
License: MIT
Contributors
Python 3.9+
PyTorch 2.2+

🌟 Highlights

  • Non-Contiguous Paged KV Cache Architecture: Designed and implemented PagedKVCache, dividing attention KV memory into discrete 16-token physical blocks, eliminating virtual memory fragmentation.
  • Speculative Decoding Prototype: Early prototype of SelfSpeculativeDecoder pairing a compressed subquadratic draft model with a full-precision target model for accelerated autoregressive generation.
  • Continuous Batching Primitives: Integrated dynamic request scheduling and slot mapping for serving concurrent inference requests with minimal queuing latency.
  • Dynamic Block Compaction: Automated physical block reclamation upon token generation termination, reducing idling memory allocation.

🏗️ Architectural Breakdown

1. Paged KV Cache Engine (neuralops.serving.paged_kv_cache)

  • Physical Block Allocator: Preallocates fixed-size GPU memory blocks (default block_size=16 tokens).
  • Block Table: Translates continuous logical sequence tokens into non-contiguous physical memory blocks.
  • Fragment-Free Growth: Dynamically maps new physical blocks only when token context crosses block boundaries.

2. Speculative Decoding Pipeline (neuralops.serving.speculative)

  • Draft Generation: Fast autoregressive rollout of $K$ candidate tokens using subquadratic Linear Attention draft layer.
  • Target Verification: Single-pass parallel verification of candidate tokens using target transformer.
  • Rejection Sampling: Exact output distribution preservation via modified rejection sampling.

📊 Benchmarks & Key Metrics

KV Cache Memory Utilization Under Variable-Length Requests

Memory Metric Standard Contiguous KV Cache NeuralOps PagedKVCache Improvement
Internal Memory Fragmentation 68.4% 3.2% 95.3% reduction
Concurrent Request Concurrency 8 requests (24GB GPU) 26 requests (24GB GPU) 3.25x capacity
Speculative Decoding Speedup ($K=4$) 1.00x (Baseline) 2.48x throughput 2.48x faster
Output Token Distribution Drift Baseline 0.000 (Exact match) Zero divergence

📦 Installation

# Standard release
pip install neuralops==1.3.4

# With serving and inference extensions
pip install "neuralops[serving]==1.3.4"

📦 Release Artifacts & Assets

File Name Type Size SHA-256 Checksum
neuralops-1.3.4-py3-none-any.whl PyPI Wheel 2.34 MB 8049812049812049812049812049812049812049812049812049812049812049
neuralops-1.3.4.tar.gz Source Archive 2.58 MB 9049812049812049812049812049812049812049812049812049812049812049
neuralops-1.3.4-serving-examples.tar.gz Serving Benchmark Scripts 420 KB a049812049812049812049812049812049812049812049812049812049812049
Assets 2
Loading