gllm

A lightweight vLLM-style inference engine built from scratch, forked from nano-vllm and extended with W8A8 int8 quantization, Blackwell-compatible attention, and several correctness/perf fixes.

Key Features

• 🧮 W8A8 int8 quantization — offline convert.py + inference-time support (per-output-channel weight scales, dynamic per-token activation scales, RTN).
• ⚡ Blackwell-ready attention — works with flash_attn_4_sm120 (sm120); decode path uses a hand-written, tile-parallel Triton gather_kvcache kernel.
• 📦 Paged KV cache with prefix caching — shared paged block manager with content-based hashing, plus chunked prefill.
• 🔀 Tensor parallelism — TP-aware weight loaders for every linear layer (column / row / merged-column / QKV), including the quantized variants.
• 🧪 Tooling — bench_quant.py (multi-scenario throughput matrix with crash retry), tests/test_w8a8.py (kernel + TP loader unit tests).

The engine lives in a single Python package (~1,900 lines including the quantization kernels) that mirrors the nano-vLLM code layout, so it stays easy to read and extend.

---

Differences from nano-vllm

gllm shares the same architecture and file layout as nano-vllm. The differences are summarized below and detailed in the sections that follow.

Aspect	nano-vllm	gllm
Weight quantization	None (bf16 only)	W8A8 int8 (gllm/layers/w8a8.py, convert.py, quant_config.json)
Attention backend	import flash_attn, flash_attn_with_kvcache for decode	Try/except for flash_attn_4_sm120; Triton gather_kvcache + flash_attn_varlen_func for decode
torch.cumsum	default (broken int64 under torch 2.13)	explicit dtype=torch.int32
dtype loading	hf_config.dtype only	handles torch_dtype string / object
TP + quantization	weight sharding only	weight and per-channel scale sharded together
Config	—	Config.quantized auto-detected from quant_config.json
Benchmark / tests	bench.py	bench_quant.py (matrix + retry), tests/test_w8a8.py

1. W8A8 int8 quantization

The biggest addition is a full quantization pipeline:

• gllm/layers/w8a8.py — Triton kernels:
  • quantize_activation — dynamic per-token activation quantization (bf16 → int8 + scale).
  • w8a8_gemm — int8 × int8 GEMM with int32 accumulation and a fused epilogue (x_scale * w_scale, optional bias), emitting bf16.
  • A TP-consistent activation quantizer (quantize_activation_tp) that reduces per-token absmax across ranks for row-parallel layers.
• gllm/utils/convert.py + convert.py — offline converter that produces a new checkpoint directory with int8 weights + fp32 {param}_weight_scale tensors and a quant_config.json marker, keeping the original HF layout (so the existing packed_modules_mapping weight loaders keep working).
• gllm/layers/linear.py — every parallel linear got a quantized variant: weight becomes int8, weight_scale (fp32, per output channel) is added, and forward dispatches to the Triton path when quantized (bf16 F.linear fallback is preserved). Row-parallel layers use the TP-consistent activation scale before all_reduce.
• gllm/models/qwen3.py / config.py / model_runner.py — thread a quantized flag through the model builder; Config reads quant_config.json to enable the W8A8 path automatically.
• lm_head / embeddings / norms are intentionally not quantized (precision sensitive / lookup-only).

2. Attention backend for Blackwell

nano-vllm imports flash_attn directly and uses flash_attn_with_kvcache for paged decode. gllm instead:

• Tries flash_attn_4_sm120 first (Blackwell), falling back to flash_attn.
• Implements a Triton gather_kvcache kernel (tile-parallel, TILE=32 rows per CTA) that materializes the paged KV cache into a contiguous buffer for flash_attn_varlen_func. This is 3–8× faster than the original whole-page-per-CTA kernel and roughly doubles decode throughput.
• Fixes out-of-bounds reads in that kernel (grid sizing and block < 0 padding guards).

3. Correctness fixes

• torch.cumsum dtype — torch ≥ 2.13 returns int64 for int32 input; the decode path passes this to flash-attn, which requires int32. Fixed with explicit dtype=torch.int32.
• flash-attn input contiguity — the prefill v is a non-contiguous qkv.split view; inputs are now .contiguous() before calling flash-attn.
• w8a8_gemm padding masks — padded M/N lanes were reading out of bounds for small decode batches; both M and N masks are now applied.

4. Benchmark & test tooling

• bench_quant.py — runs a bf16/w8a8 × batch-size matrix (each scenario in a fresh subprocess, repeated and auto-retried), reporting prefill / decode throughput and peak GPU memory.
• tests/test_w8a8.py — Triton kernel equivalence vs F.linear, activation quantizer correctness, CUDA-graph replay, and TP sharding of int8 weights + scales.

---

Installation

pip install -e .

Requires PyTorch ≥ 2.12, Triton ≥ 3.0, flash-attn (either flash_attn or flash_attn_4_sm120), safetensors, transformers.

Quantize a model (W8A8)

# Produce a quantized checkpoint directory next to the original
python convert.py --input ~/models/Qwen3-0.6B --output ~/models/Qwen3-0.6B-w8a8

The output directory keeps config.json / tokenizer files, adds int8 weights + fp32 scales, and a quant_config.json:

{"quant_method": "w8a8", "weight_dtype": "int8", "scale_dtype": "fp32"}

Quick Start

The API mirrors vLLM / nano-vllm:

from gllm import LLM, SamplingParams

# Load the quantized checkpoint (auto-detects quant_config.json)
llm = LLM("~/models/Qwen3-0.6B-w8a8", enforce_eager=True)

sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
outputs = llm.generate(["Hello, Nano-vLLM."], sampling_params)
print(outputs[0]["text"])

Running an unquantized (bf16) directory works too — the same LLM class falls back to F.linear automatically.

Quantization Design

• Scheme: symmetric W8A8.
• Weight scale: per output channel, amax / 127, computed from the original weights at conversion time (round-to-nearest, clamped to [-127, 127]).
• Activation scale: dynamic per token, computed at runtime (amax / 127).
• No calibration data required.
• Coverage: all 4 linear families (QKV / gate-up / down / o); lm_head, embeddings and norms stay bf16.
• TP: scales are sharded consistently with the weights, so tensor-parallel inference works with quantized models.

Benchmark

Test configuration

• Hardware: NVIDIA RTX PRO 5000 72GB (Blackwell, sm120)
• Software: PyTorch 2.13.0+cu130, Triton 3.7.1
• Model: Qwen3-0.6B (hidden=1024, 28 layers, 16 heads, 8 KV heads, head_dim=128)
• Load: input 100–1024 tokens, output 100–1024 tokens, random; ignore_eos=True
• Mode: eager (enforce_eager=True); each scenario averaged over 2 runs

Checkpoint size (weights only)

checkpoint	size
bf16	1.40 GiB
w8a8 (int8)	0.99 GiB

Throughput (6 scenarios)

Scenario	bf16 prefill (tok/s)	w8a8 prefill (tok/s)	bf16 decode (tok/s)	w8a8 decode (tok/s)	decode Δ
eager bs=8	77,160	28,029	208	162	-22%
eager bs=32	79,846	51,241	755	560	-26%
eager bs=128	70,930	56,320	1,889	1,551	-18%

Peak GPU memory is ~64.5 GiB for all scenarios (gpu_memory_utilization=0.9; KV cache dominates and is not quantized).

Accuracy (quantized vs bf16)

metric	result	threshold
logits max-abs-diff	1.39	< 2.0
greedy top-1 agreement	100% (both predict Paris)	> 90%

Interpretation

• End-to-end W8A8 is not faster on this small model (decode 18–26% slower). On a 0.6B model the int8 GEMM overhead (quantize kernel + extra launches, 28 layers × 4 linears) outweighs the tensor-core win, and cuBLAS bf16 is very fast at tiny decode batch sizes. At the pure-GEMM level W8A8 is ~1.3× faster for large M (prefill-like shapes), but the benefit is diluted end-to-end.
• The clear win is -29% checkpoint size (memory / bandwidth), which pays off most for memory-bound long-context serving or larger models.

Known Limitations

• CUDA graph path (enforce_eager=False) is currently broken in this repo (pre-existing): gather_kvcache allocates via torch.empty during graph capture, which is disallowed. Use enforce_eager=True.
• flash_attn_4_sm120 (v0.1.0) still has an occasional (~10%) illegal memory access on large prefills (CUTLASS-kernel internal, environment/library issue, not introduced by quantization). bench_quant.py retries each scenario to work around it; PYTORCH_NO_CUDA_MEMORY_CACHING=1 avoids it but is very slow.
• Quantized inference does not yet support the CUDA-graph decode path.

Files

• Added: convert.py, gllm/layers/w8a8.py, gllm/utils/convert.py, tests/test_w8a8.py, bench_quant.py, bench_results.md
• Modified: gllm/layers/linear.py (W8A8 variants), gllm/layers/attention.py (gather / cumsum fixes), gllm/models/qwen3.py (quantized flag), gllm/config.py, gllm/engine/model_runner.py, gllm/engine/llm_engine.py