License notice This project is source-available under the PolyForm Noncommercial License 1.0.0 — free for personal, research, teaching, and internal-evaluation use.

Any commercial use, redistribution as part of a commercial product, or paid hosted deployment requires a separate commercial license. See COMMERCIAL.md or contact https://github.com/AnubhabBanerjee.

Unauthorized commercial use or redistribution is a violation of the license terms.

---

Read the full deep-dive here: https://towardsdatascience.com/i-built-a-c-backend-so-my-gpu-would-stop-eating-air/

---

---

🚀 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.

🎯 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. Packed bins are 1-D token ribbons plus cu_seqlens cut-points. VarlenModelWrapper now routes those cut-points into flash_attn_varlen_func, so documents in the same bin do not attend to each other. Load the Hugging Face model with attn_implementation="flash_attention_2" (see main_working_file.py / example_runs/optimized_run.py).

Benchmark note: The tables below were measured with the older packed path + SDPA (standard PyTorch attention). In that configuration cu_seqlens reset position IDs only; attention still treated the bin as one sequence. Those numbers are throughput/memory, not a claim of identical logits vs the padded baseline. Re-run optimized_run.py after this change to measure the FA2-varlen stack.

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, SM 6.1) | Model: SmolLM2-360M-Instruct Stack: WarpGroup packed 1-D bins + SDPA (not FlashAttention-2)

Pascal (SM 6.1) cannot run FlashAttention-2. This run therefore never called flash_attn_varlen_func. Packed documents shared one SDPA sequence, so later documents could attend to earlier ones. 5.89× is a throughput/memory number on that broken attention configuration, not a correct-varlen result. The current attn_implementation="flash_attention_2" path will not run on this GPU.

By packing continuous 1D tensors, WarpGroup still avoided pad tokens and used less VRAM than the rectangular baseline — that part of the table is the padding story, not a correctness claim.

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 (SDPA, cross-doc attention; not FA2)
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.
6. Phase 5: Varlen attention — VarlenModelWrapper rebuilds per-document position_ids from cu_seqlens and patches FlashAttention-2 so each layer calls flash_attn_varlen_func with those cuts. Hugging Face model.forward() still does not take cu_seqlens; the wrapper is the bridge.

🛠️ 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 (flash-attn)	Variable-length attention via flash_attn_varlen_func (requires attn_implementation="flash_attention_2").
Hugging Face Transformers	Base model loading and tokenizer. VarlenModelWrapper injects cu_seqlens into FA2 layers.

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. Packed inference requires flash-attn and attn_implementation="flash_attention_2" so cu_seqlens reach flash_attn_varlen_func.

⚙️ 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 autotuner + VarlenModelWrapper (FA2 varlen + position reset)
├── 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 stack (FFD + zero-copy + FA2 varlen)
    ├── 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.