Building a Documentation Generator with vLLM: Complete Guide

Building a Documentation Generator with vLLM: Complete Guide

Documentation is often the most neglected part of the software development lifecycle. Developers write code, ship features, and move on—leaving behind sparse READMEs, outdated docstrings, and confused future maintainers. With the rise of high-throughput inference engines like vLLM, it's now practical to build a documentation generator that reads your codebase and produces accurate, context-aware documentation at scale.

This tutorial walks you through building a complete documentation generator powered by vLLM. We'll cover what vLLM is, why it's ideal for this use case, how to set it up, and how to build a production-ready generator that parses source files, batches prompts efficiently, and writes structured Markdown documentation.

What Is vLLM?

vLLM is an open-source inference engine for Large Language Models (LLMs) developed at UC Berkeley. It is designed to deliver state-of-the-art throughput and memory efficiency through two key innovations:

• PagedAttention: A novel attention algorithm that manages KV cache memory like an operating system manages virtual memory, drastically reducing waste.
• Continuous batching: Dynamically inserts and completes requests in a batch, keeping the GPU saturated instead of waiting for the slowest request in a static batch.

For a documentation generator, these features matter because you'll often process dozens or hundreds of source files. vLLM lets you batch these requests efficiently, cutting inference time and cost dramatically compared to naive sequential generation.

Why Use vLLM for Documentation Generation?

Generating documentation for a real codebase is a batch inference problem at its core. Each source file (or function, or class) becomes a prompt, and you want to process as many as possible in parallel. Here's why vLLM stands out:

• High throughput: vLLM can be 10–20x faster than naive HuggingFace Transformers for batched workloads.
• OpenAI-compatible API: You can run vLLM as a server and use the standard openai Python client, making it easy to swap models.
• Local and private: Your code never leaves your machine. This is critical for proprietary codebases.
• Model flexibility: vLLM supports most popular open-weight models including Llama, Mistral, Qwen, and CodeLlama families.

Prerequisites and Setup

Before we start coding, make sure you have the following:

• A machine with an NVIDIA GPU (at least 16GB VRAM for a 7B model, more for larger models).
• Python 3.10 or newer.
• Docker (optional, but recommended for running the vLLM server).
• A codebase you want to document.

Installing vLLM

The simplest way to install vLLM is via pip. Create a virtual environment and install the package:

python -m venv venv
source venv/bin/activate
pip install vllm

If you plan to use the OpenAI-compatible client (recommended), also install the client library:

pip install openai

Starting the vLLM Server

For most real-world use cases, running vLLM as a server is the cleanest approach. It exposes an OpenAI-compatible REST API that your documentation generator can call. Start the server with a capable code-aware model:

vllm serve Qwen/Qwen2.5-Coder-7B-Instruct \
  --port 8000 \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.9

This command downloads the model (if not cached), loads it into GPU memory, and starts serving on port 8000. The --max-model-len flag controls the maximum context window, which is important because source files can be long.

Architecture of the Documentation Generator

Our generator will follow a simple pipeline:

• File discovery: Walk the codebase and collect source files.
• Chunking: Split large files into manageable chunks (functions, classes, or fixed-size blocks).
• Prompt construction: Build a structured prompt for each chunk.
• Batched inference: Send all prompts to vLLM in batches.
• Output assembly: Merge generated documentation into Markdown files.

Let's build each component.

Step 1: File Discovery and Chunking

We'll use Python's pathlib and ast modules to discover Python files and split them into logical units. The approach below parses each file into an Abstract Syntax Tree and extracts top-level functions and classes.

import ast
from pathlib import Path
from dataclasses import dataclass


@dataclass
class CodeChunk:
    file_path: str
    name: str
    kind: str  # "function" or "class"
    source: str
    start_line: int


def discover_files(root: str, extensions: tuple = (".py",)) -> list[Path]:
    """Walk a directory tree and return matching source files."""
    root_path = Path(root)
    files = []
    for path in root_path.rglob("*"):
        if path.is_file() and path.suffix in extensions:
            # Skip common non-source directories
            if any(part in {"venv", ".git", "__pycache__", "node_modules"} 
                   for part in path.parts):
                continue
            files.append(path)
    return files


def chunk_python_file(file_path: Path) -> list[CodeChunk]:
    """Parse a Python file and extract top-level functions and classes."""
    source_text = file_path.read_text(encoding="utf-8")
    try:
        tree = ast.parse(source_text)
    except SyntaxError:
        return []

    chunks = []
    lines = source_text.splitlines(keepends=True)

    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            kind = "function"
        elif isinstance(node, ast.ClassDef):
            kind = "class"
        else:
            continue

        start = node.lineno - 1
        end = node.end_lineno if node.end_lineno else len(lines)
        snippet = "".join(lines[start:end])

        chunks.append(CodeChunk(
            file_path=str(file_path),
            name=node.name,
            kind=kind,
            source=snippet,
            start_line=node.lineno,
        ))

    return chunks

This gives us clean, semantically meaningful chunks rather than arbitrary line-based splits. Each chunk contains a single function or class, which helps the model produce focused documentation.

Step 2: Prompt Construction

The quality of generated documentation depends heavily on prompt design. We'll use a system prompt that defines the assistant's role and a user prompt that provides the code and asks for structured output.

PROMPT_TEMPLATE = """You are a senior software engineer writing documentation.

Below is a {kind} named `{name}` from the file `{file_path}`.

Write clear, professional documentation in Markdown for this code. Include:
1. A one-sentence summary of what it does.
2. A detailed description of its behavior and purpose.
3. A table of parameters or attributes (if applicable).
4. A usage example.
5. Any important notes, edge cases, or caveats.

Use a `##` heading for the title. Do not include the original source code.

### Source code:
python
{source}
"""


def build_prompt(chunk: CodeChunk) -> str:
    return PROMPT_TEMPLATE.format(
        kind=chunk.kind,
        name=chunk.name,
        file_path=chunk.file_path,
        source=chunk.source,
    )

Notice that we explicitly ask the model not to repeat the source code. This keeps the output concise and focused on documentation only.

Step 3: Batched Inference with vLLM

Now we connect to the vLLM server and send our prompts in batches. Because vLLM handles continuous batching internally, we can send a large number of requests and let the server manage scheduling. However, we'll still chunk our requests client-side to avoid overwhelming memory.

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="dummy",  # vLLM doesn't require a real key by default
)

MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"


def generate_docs(prompts: list[str], batch_size: int = 16) -> list[str]:
    """Send prompts to vLLM in batches and return generated text."""
    results = []

    for i in range(0, len(prompts), batch_size):
        batch = prompts[i:i + batch_size]
        print(f"Processing batch {i // batch_size + 1} "
              f"({len(batch)} prompts)...")

        responses = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "user", "content": prompt}
                for prompt in batch
            ],
            temperature=0.2,
            max_tokens=1024,
        )

        for choice in responses.choices:
            results.append(choice.message.content)

    return results

We use a low temperature (0.2) because documentation should be factual and deterministic. The max_tokens limit prevents runaway generations from consuming GPU time.

Step 4: Assembling the Output

After generating documentation for each chunk, we assemble everything into a structured set of Markdown files. A common approach is to mirror the source directory structure under an docs/ folder.

def assemble_docs(chunks: list[CodeChunk], docs: list[str], output_dir: str):
    """Write generated documentation to Markdown files."""
    output_root = Path(output_dir)
    output_root.mkdir(parents=True, exist_ok=True)

    # Group chunks by source file
    by_file: dict[str, list[tuple[CodeChunk, str]]] = {}
    for chunk, doc in zip(chunks, docs):
        by_file.setdefault(chunk.file_path, []).append((chunk, doc))

    for file_path, entries in by_file.items():
        # Build output path: src/utils/helpers.py -> docs/utils/helpers.md
        rel = Path(file_path).with_suffix(".md")
        out_path = output_root / rel.name
        out_path.parent.mkdir(parents=True, exist_ok=True)

        content = [f"# Documentation: {rel}\n"]

        # Sort by line number to preserve source order
        entries.sort(key=lambda x: x[0].start_line)

        for chunk, doc in entries:
            content.append(f"\n\n")
            content.append(doc)

        out_path.write_text("\n".join(content), encoding="utf-8")
        print(f"Wrote {out_path}")

The HTML comment at the top of each section records the source location, making it easy to trace documentation back to the original code.

Step 5: Putting It All Together

Now we combine all components into a single entry point:

def main():
    import sys

    if len(sys.argv) < 2:
        print("Usage: python docgen.py <source_dir> [output_dir]")
        sys.exit(1)

    source_dir = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "docs"

    # Step 1: Discover and chunk files
    print(f"Scanning {source_dir}...")
    files = discover_files(source_dir)
    print(f"Found {len(files)} source files.")

    all_chunks = []
    for f in files:
        all_chunks.extend(chunk_python_file(f))

    print(f"Extracted {len(all_chunks)} code chunks.")

    if not all_chunks:
        print("No chunks found. Exiting.")
        return

    # Step 2: Build prompts
    prompts = [build_prompt(c) for c in all_chunks]

    # Step 3: Generate documentation
    print("Generating documentation via vLLM...")
    docs = generate_docs(prompts, batch_size=16)

    # Step 4: Assemble output
    assemble_docs(all_chunks, docs, output_dir)
    print(f"Done! Documentation written to {output_dir}/")


if __name__ == "__main__":
    main()

Run the generator against any Python project:

python docgen.py ./my_project ./generated_docs

Best Practices

Choose the Right Model

Code documentation requires strong code understanding. Models like Qwen2.5-Coder-7B-Instruct, CodeLlama-13B-Instruct, or DeepSeek-Coder-V2-Lite-Instruct are good choices. Larger models produce higher-quality prose but require more VRAM and are slower. Start with a 7B model and scale up if quality is insufficient.

Handle Large Files Gracefully

Some files may exceed your model's context window. Add a length check before building prompts and either skip them or split them further:

MAX_SOURCE_CHARS = 6000  # Roughly 1500 tokens of code

def build_prompt(chunk: CodeChunk) -> str | None:
    if len(chunk.source) > MAX_SOURCE_CHARS:
        print(f"Skipping {chunk.name}: source too long ({len(chunk.source)} chars)")
        return None
    return PROMPT_TEMPLATE.format(
        kind=chunk.kind,
        name=chunk.name,
        file_path=chunk.file_path,
        source=chunk.source,
    )

Use Deterministic Settings

Set temperature=0 or a very low value for reproducible output. Documentation should be stable across runs so that diffs in version control are meaningful. If you regenerate docs after a small code change, only the affected sections should change.

Include Context from Surrounding Code

A function in isolation can be ambiguous. Consider including imports, module-level docstrings, or related class definitions as context. You can modify the chunker to capture the file header:

def extract_file_header(file_path: Path, max_lines: int = 30) -> str:
    lines = file_path.read_text(encoding="utf-8").splitlines()
    header_lines = []
    for line in lines[:max_lines]:
        header_lines.append(line)
        if line.strip() and not line.startswith(("#", "import", "from")):
            break
    return "\n".join(header_lines)

Then include this header in the prompt so the model understands the module's imports and intent.

Validate and Post-Process Output

LLMs occasionally produce hallucinated APIs or incorrect parameter names. For production use, add a validation step that checks whether parameter names in the generated docs match the actual function signature:

import re

def validate_params(chunk: CodeChunk, doc: str) -> list[str]:
    """Check that documented parameters match the actual signature."""
    # Extract actual parameters from AST
    tree = ast.parse(chunk.source)
    func_node = tree.body[0]
    if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
        return []
    
    actual_params = {
        arg.arg for arg in func_node.args.args
        if arg.arg != "self"
    }
    
    # Extract parameter names mentioned in the doc
    mentioned = set(re.findall(r"`(\w+)`", doc))
    
    missing = actual_params - mentioned
    extra = mentioned - actual_params - {"self", "cls"}
    
    issues = []
    if missing:
        issues.append(f"Missing parameters: {missing}")
    if extra:
        issues.append(f"Unknown parameters: {extra}")
    return issues

Cache Results to Avoid Regeneration

Generating documentation for a large codebase takes time. Implement a simple hash-based cache so that only changed files are reprocessed:

import hashlib
import json
from pathlib import Path

def file_hash(path: Path) -> str:
    return hashlib.md5(path.read_bytes()).hexdigest()

def load_cache(cache_path: str) -> dict:
    p = Path(cache_path)
    if p.exists():
        return json.loads(p.read_text())
    return {}

def save_cache(cache: dict, cache_path: str):
    Path(cache_path).write_text(json.dumps(cache, indent=2))

Store a mapping of file_path -> hash and skip files whose hash hasn't changed since the last run.

Monitor vLLM Server Health

When processing hundreds of chunks, the vLLM server can become a bottleneck. Monitor GPU utilization with nvidia-smi and adjust batch_size accordingly. If you see the server returning errors, reduce the batch size or add retry logic with exponential backoff.

Extending the Generator

Once the basic pipeline works, you can extend it in several directions:

• Multi-language support: Add chunkers for JavaScript (using a parser like esprima), Go, or Rust. The prompt template and inference logic remain the same.
• Cross-reference linking: After generating docs, run a second pass that detects references to other functions/classes and inserts Markdown links.
• README generation: Feed the list of all generated module docs into a final prompt that produces a top-level README summarizing the project.
• CI/CD integration: Run the generator in a GitHub Actions workflow on every pull request, committing updated docs automatically.

Conclusion

Building a documentation generator with vLLM is a practical and powerful way to keep your codebase documented without manual effort. By combining vLLM's high-throughput inference with a clean chunking and prompting pipeline, you can process entire projects in minutes rather than hours. The key to good results lies in thoughtful prompt design, deterministic generation settings, and post-generation validation to catch hallucinations. Start with a 7B code model, iterate on your prompts, and gradually add caching, multi-language support, and CI integration as your needs grow. With this foundation, you have a flexible system that turns source code into maintainable, professional documentation at scale.