Copyright (c) 2026 Anubhab Banerjee (AnubhabBanerjee/WarpGroup-backend) All rights reserved. No part of this repository may be used, redistributed, or modified in any form or by any means without the prior written permission of the author.

---

🚀 WarpGroup-Backend: VRAM-Aware Asynchronous Sequence Packing & Zero-Copy Inference

Zero-padding, asynchronous First-Fit Decreasing (FFD) bin packing for extreme-context LLM workloads (e.g., AI patent evaluation). Bypasses the GIL and CPU memory bottlenecks by streaming tokenized sequences into a C++ backend, packing them dynamically based on empirical VRAM hardware limits, and passing them to FlashAttention-2 via zero-copy pinned memory (cudaHostAlloc).

This repo is a high-performance inference infrastructure slice: Python handles dynamic VRAM discovery and generic stream tokenization, while a compiled C++ backend (exposed via PyBind11) manages thread-safe sequence queuing, 16-token Tensor Core alignment, and exact hardware-bounded bin packing without ever duplicating data in host RAM.

For a complete explanation and tutorial on this repository, please read this article: https://towardsdatascience.com/i-built-a-c-backend-so-my-gpu-would-stop-eating-air/

🎯 Why VRAM-aware packing matters

In extreme-context LLM inference, variable-length documents cause massive padding waste or fatal Out-Of-Memory (OOM) errors with standard item-count batching. Kubernetes or native PyTorch schedulers do not protect against dynamic attention workspace bloat.

This repository implements a systems-level paradigm shift: it transitions from "Batch Size" batching to "VRAM Capacity" packing. By micro-padding to hardware boundaries and dynamically grouping independent sequences into a strictly enforced, hardware-defined byte limit, it guarantees zero OOMs and ensures every forward pass utilizes maximum GPU silicon.

📈 Performance benchmarks

WarpGroup was evaluated across entry-level and production hardware on highly variable text corpora. It seamlessly integrates with flash_attn_varlen_func to saturate Tensor Cores without the drag of padded tokens.

1. Extreme variance stress test (the "real world" distribution)

Hardware: NVIDIA H100 (80 GB) | Model: Qwen2.5-7B-Instruct Stack: WarpGroup tuned stack + SDPA (Standard PyTorch attention) Dataset: 400 PDFs (high-variance skew: 45–130-word documents interleaved with 1820–2000-word documents)

Standard batching collapses when faced with high-variance document lengths. In this stress test, the baseline Hugging Face pipeline was forced to pad over 818,000 tokens just to maintain rectangular tensor shapes, inflating its dynamic memory usage and cutting throughput in half.

WarpGroup effortlessly absorbed the variance, using 1D continuous tensors to achieve a 2.08× throughput multiplier and a ~62% reduction in dynamic memory overhead.

Metric	Baseline (HF)	WarpGroup	Improvement
Padding overhead	48.41%	0.55%	47.9 pp absolute reduction
Throughput	14,713 tok/s	30,672 tok/s	2.08× higher
Peak VRAM	19.88 GB	16.50 GB	17% lower (3.38 GB saved)
Dynamic VRAM (est.)*	~5.38 GB	~2.00 GB	~62% lower dynamic memory
Wall clock	28.69 s	13.76 s	2.08× faster

*Architecture note on dynamic VRAM: A 7B-parameter model in bf16/fp16 requires ~14.5–15 GB of static VRAM for weights. The dynamic memory (activations / KV cache) is where the padding penalty occurs. WarpGroup processed the exact same sequence of real tokens while shedding over 3 GB of wasted dynamic allocation.

2. Production scaling (uniform variable-length)

Dataset: 300 PDFs (uniform distribution: 50–1900 words)

Even on a less adversarial, uniformly distributed dataset, WarpGroup fully saturates the Hopper Tensor Cores without the drag of padded tokens, yielding a 70% increase in useful-token throughput over the baseline.

Metric	Baseline (HF)	WarpGroup	Improvement
Padding overhead	36.20%	0.67%	35.5 pp absolute reduction
Throughput	18,047 tok/s	30,700 tok/s	1.70× higher
Wall clock	17.86 s	10.50 s	1.70× faster

3. Entry-level hardware (the padding annihilation)

Hardware: NVIDIA GeForce GTX 1080 (8 GB) | Model: SmolLM2-360M-Instruct

By packing continuous 1D tensors, WarpGroup drastically rescues throughput on older or memory-constrained hardware where standard batching wastes extreme amounts of compute.

Metric	Baseline (HF)	WarpGroup	Improvement
Padding overhead	41.13%	0.00%	Baseline padding eliminated
Throughput	405 tok/s	2,387 tok/s	5.89× higher
Peak VRAM	2.85 GB	1.86 GB	35% lower

4. OOM prevention (unbounded lengths)

When MAX_LEN constraints are removed, standard batching attempts to allocate memory for the theoretical maximum grid (batch_size × longest_sequence), rapidly causing CUDA Out-Of-Memory errors on variable text.

• Baseline: Crashed (torch.OutOfMemoryError: Tried to allocate 30.00 GiB).
• WarpGroup: Completed successfully (Peak VRAM: 3.60 GB). The Phase-0 autotune probes the GPU at startup and locks a strict hardware-aligned token budget; the bin packer rejects any forward pass that would exceed that budget.

Reproducing the benchmarks

python example_runs/run_vram_contrast_benchmark.py \
  --model example_models/Qwen2.5-7B-Instruct \
  --baseline-batch-size 8

python example_runs/random_pdf_generator.py --count 300 \
  --tokens-min 50 --tokens-max 1900 --seed 42
python example_runs/baseline_run.py  --pdf-dir example_runs/data/ \
  --model example_models/Qwen2.5-7B-Instruct --batch-size 4
python example_runs/optimized_run.py --pdf-dir example_runs/data/ \
  --model example_models/Qwen2.5-7B-Instruct

Raw JSON for every run is written under example_runs/results/ (and copied into example_runs/plots/<run_id>/ for archival). The §1 manifest is example_runs/results/vram_contrast_manifest.json.

🧠 System architecture

The pipeline is a decoupled autotune → ingest → pack → execute graph—I/O is isolated in Python, while compute and memory mapping run concurrently in C++:

1. Phase 0: Hardware Autotuning — determine_vram_capacity empirically probes the GPU with synthetic sequences to find the physical VRAM limit for your specific model's attention workspace.
2. Phase 1: Generic Ingestion — Python generators yield strings (e.g., from JSONL), tokenize them to flat integers, and stream them via submit_sequence across the PyBind11 boundary.
3. Phase 2: Async C++ Dispatch & Alignment — async_dispatcher.cpp catches tokens in a std::deque outside the Python GIL. Sequences are micro-padded to 16-token intervals to prevent Tensor Core stalling.
4. Phase 3: FFD Bin Packing — A background thread sorts the queue and packs sequences into a pinned memory pool (cudaHostAlloc) using a First-Fit Decreasing algorithm up to the empirical VRAM limit.
5. Phase 4: Zero-Copy Wrapping — torch::from_blob wraps the C++ memory in a PyTorch metadata shell. The GPU's DMA controller pulls the unpadded sequences directly across the PCIe bus for FlashAttention-2 execution.

🛠️ Stack & core backend

Layer	Role
C++17 / PyBind11	Thread-safe queueing, memory pooling, and GIL-free background execution.
CUDA Runtime API	Explicit page-locked host memory allocation (cudaHostAlloc).
PyTorch 2.5+	Tensor metadata wrapping, model weights, and graph execution.
FlashAttention-2	Variable-length, padding-free exact attention (flash_attn_varlen_func).
Hugging Face Transformers	Base model loading and tokenizer definitions.

Reference environment (pinned)

Item	Value
OS	Ubuntu 24.04
GPU	NVIDIA H100 (80GB) or equivalent
NVIDIA driver	535+ (CUDA 12.2+)
Python	3.12+
Compiler	GCC 9.0+ / CMake 3.18+
PyTorch	v2.5+ (cu121)

✅ Prerequisites

• Linux host for compiling the C++ backend and running the PyTorch loops.
• CUDA Toolkit installed and accessible in $PATH for compiling the C++ extensions.
• Python 3.12+ with a virtual environment.
• A target LLM supported by Hugging Face and configured for flash_attention_2.

⚙️ Installation

From the repository root, set up your Python environment and compile the C++ backend:

python3.12 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt
pip install -e .

# Compile the PyBind11 C++ backend
mkdir build && cd build
cmake ..
make -j4
cp warpgroup_backend*.so ..
cd ..

🚀 Execution

End-to-end VRAM autotuning, ingestion, and background packing:

python3 main_working_file.py

🎬 Example run

Default pipeline locations

Artifact	Path
Python Entry Point	main_working_file.py
PyTorch DataLoader Logic	streaming_dataloader.py
C++ Engine Bindings	csrc/bindings.cpp
C++ Async Queue	csrc/core/async_dispatcher.h, csrc/core/async_dispatcher.cpp
Built Extension	warpgroup_backend.cpython-312-x86_64-linux-gnu.so

📁 Project layout

├── README.md                                # Project overview + benchmarks
├── README_vram_contrast_dataset.md          # → see example_runs/  (note: this file is under example_runs/)
├── CMakeLists.txt                           # PyBind11 + CUDA build directives
├── setup.py                                 # `pip install -e .` entry point
├── requirements.txt                         # torch, transformers, flash-attn, pymupdf, ...
├── .gitignore
├── overview_image.png                       # README hero image
│
├── main_working_file.py                     # End-to-end driver
├── streaming_dataloader.py                  # Phase-0 hardware autotuner + VarlenModelWrapper
├── reader_and_tokennizer.py                 # PDF reader + tokenizer helpers
│
├── warpgroup/                               # Importable Python package
│   ├── __init__.py
│   ├── dataloader.py
│   └── reader.py
│
├── csrc/                                    # C++17 / PyBind11 backend
│   ├── bindings.cpp                         # PyBind11 Python ↔ C++ boundary
│   ├── core/
│   │   ├── async_dispatcher.cpp             # GIL-free background worker
│   │   ├── bin_packer.cpp                   # First-Fit Decreasing packing + 16-token TC alignment
│   │   └── memory_pool.cpp                  # cudaHostAlloc pinned memory + torch::from_blob handoff
│   └── include/
│       ├── async_dispatcher.h
│       ├── bin_packer.h
│       └── memory_pool.h
│
└── example_runs/                            # Reproducible benchmark harness
    ├── README_example.md
    ├── README_vram_contrast_dataset.md      # Notes for the interleaved long/short corpus
    ├── run_e2e.sh                           # CMake build + setup + baseline + optimized + report
    ├── setup_experiment.py                  # Builds + copies warpgroup_backend.so into place
    ├── random_pdf_generator.py              # Uniform variable-length synthetic PDF corpus
    ├── baseline_run.py                      # Standard Hugging Face padded-batch baseline
    ├── optimized_run.py                     # WarpGroup tuned stack (FFD + zero-copy + FA2/SDPA)
    ├── generate_report.py                   # Builds bar charts into plots/<run_id>/
    ├── build_results_docx.py                # Compiles every run into RESULTS.docx
    ├── RESULTS.docx                         # Word-format results dossier
    └── plots/                               # PNG bar charts per run (PNGs only, no JSONs in zip)
        ├── e2e_smollm_300pdf/
        │   ├── padding_overhead.png
        │   ├── throughput_tokens.png
        │   ├── time_comparison.png
        │   └── vram_comparison.png
        ├── run_A_maxlen2048_dense/{padding_overhead,throughput_tokens,time_comparison,vram_comparison}.png
        ├── run_C_maxlen2048_variable/{padding_overhead,throughput_tokens,time_comparison,vram_comparison}.png
        ├── run_D_varlen_qwen_tuned/{padding_overhead,throughput_tokens,time_comparison,vram_comparison}.png
        ├── run_vram_contrast_corpus/{padding_overhead,throughput_tokens,time_comparison,vram_comparison}.png
        └── run_vram_contrast_appendix/{padding_overhead,throughput_tokens,time_comparison,vram_comparison}.png

🛣️ Roadmap

Active and planned hardening for the engine:

• Multi-GPU Sharding — Extend the dispatcher to manage multiple C++ queues, distributing dynamically sized bins across local GPU interconnects.

🙏 Acknowledgments

Built with PyBind11, PyTorch, and the NVIDIA CUDA Toolkit to optimize single-node LLM throughput by respecting silicon-level boundaries. Architecture inspired by the constraints of high-volume, asynchronous document evaluation pipelines.