Soulmate-Halo/heterogeneous-gpu-pd-lab

Soulmate‑Halo / heterogeneous‑gpu‑pd‑lab – Dense Acceleration Lab

What it is

• A research/benchmarking codebase that explores heterogeneous GPU pipelines for large‑language‑model (LLM) inference. It shows how to split the Prefill (encoding of long prompts) onto a pair of RTX 6000D GPUs (TP2) while the Decode (generation) runs on a larger‑memory back‑stage of four DGX Spark nodes (TP4). The project calls this two‑stage “PP2” architecture and measures token‑throughput under real‑world workloads.

Key ideas & components

• PP2 pipeline – front‑stage (Prefill) on two RTX 6000D GPUs, back‑stage (Decode) on four Sparks. The front stage only needs the encoder part of a causal encoder‑decoder model, so the heavy‑weight decoder weights stay on the back stage.
• Tensor‑parallelism (TP) – TP2 for the front stage, TP4 for the back stage, coordinated via the NIXL cache‑page hand‑off.
• Frameworks – front stage runs vLLM (TP2), back stage runs SGLang (TP4/EP4). The system also uses flashinfer_cutlass, DSpark 5, and CUDA‑graph decode batches up to 32.
• Metrics – reports P‑stage batch Prefill throughput (tokens / second) and active‑Decode / end‑to‑end throughput for various concurrency levels (C16‑C32) and input lengths (8 K‑32 K tokens).
• Reproducible data – CSV tables, JSON telemetry, and markdown reports are stored in the repo’s data/ and results/ folders.

Why it matters

• Demonstrates that a modest pair of high‑end GPUs can handle the encoder side of a 8 B‑parameter model, freeing larger‑memory GPUs to focus on generation. This reduces the need for every GPU to hold the full model, cutting hardware cost for long‑context inference.
• Provides concrete throughput numbers (e.g., 16 300 tok/s Prefill and 442 tok/s aggregate Decode on a six‑GPU system) that can be used as baselines for anyone building heterogeneous LLM serving stacks.
• Shares a cost breakdown (≈ US $38 k for the six‑GPU setup) and compares it with all‑Spark or all‑H20 configurations, helping practitioners evaluate trade‑offs.

Typical use‑case

• Organizations that need to serve very long prompts (up to 32 K tokens) with low latency but cannot afford a full‑model replica on every GPU. By pairing a small‑VRAM, high‑throughput card (RTX 6000D) with a large‑memory server (DGX Spark), they can achieve high Prefill rates while keeping generation quality.

Maturity

• The repository contains detailed benchmark reports up to V8 (seq‑32), with historical data from V1‑V7. It is a research/engineering artifact rather than a polished library; deployment scripts, patches, and the exact layer‑allocation policy are kept private. Nonetheless, the code and data are sufficient to reproduce the reported numbers on similar hardware.

Getting started (high‑level)

1. Hardware – two RTX 6000D GPUs + four DGX Spark nodes (or equivalent with comparable memory/compute). Ensure NVLink/NVSwitch connectivity for the back‑stage cluster.
2. Software stack – install vLLM (for the front stage) and SGLang (for the back stage). Enable flashinfer_cutlass and set DSpark=5.
3. Configuration – follow the V8 settings described in the README:
   • Front‑stage TP2, back‑stage TP4/EP4
   • Decode CUDA‑graph max_bs=32
   • Chunk size 2048, tail‑replay 256
   • Use the provided model YAML for DeepSeek‑V4.1‑Flash (or other CED models).
4. Run – launch the front‑stage service, then the back‑stage service, and submit batch requests (e.g., 32 K‑token prompts, 1024‑token budget). Use the CSV/JSON files in data/ as reference for expected telemetry.
5. Validate – check that all requests succeed (the repo reports 216/216 success for the V8 run) and compare your throughput against the tables in the README.

Where to look for more info

• results/ – markdown deep‑dives for each model/hardware combo.
• data/ – raw CSV and JSON evidence files.
• The linked external docs: model description on HuggingFace, vLLM recipe YAML, and the SGLang repository.

Bottom line heterogeneous-gpu-pd-lab is a concrete, performance‑focused project that shows how to split LLM inference across GPUs with different memory/compute characteristics, delivering record Prefill speeds for long‑context models while keeping generation costs manageable.