00_README.md

How the OpenBMB MiniCPM5 inference speedup was found

I asked GPT Sol to optimize inference for openbmb/MiniCPM5-1B on an NVIDIA RTX 3060.

The goal was not simply to produce a larger tokens-per-second number. A candidate only counted as an improvement if it produced the same generated token IDs as the frozen baseline.

The companion file 01_core_code_snippets.py shows the central implementation ideas: stable CUDA Graph buffers, fixed-shape replay, combining adjacent operations, grouping Q/K/V work, and falling back to eager execution for other shapes. It is intentionally illustrative rather than a standalone drop-in script.

Starting point

The frozen eager baseline was:

• Model: openbmb/MiniCPM5-1B
• Hardware: NVIDIA RTX 3060
• Precision: BF16
• Workload: batch-one greedy generation
• Test set: six held-out prompts
• Measurement: two warmups, five measured repetitions, 64 generated tokens
• Baseline: 42.95 decode tokens/second

Every serious candidate was tested in a fresh process. The output token trace was compared with the baseline, and promising candidates received a separate confirmation run.

Why torch.compile was rejected

One automatic compilation configuration appeared to reach 120.9 tokens/second. It was not accepted: generated token IDs differed from the baseline in 15 of 30 comparisons.

That result was faster but incorrect for this experiment. The harness therefore treated it as a rejected candidate rather than a real speedup.

Profiling showed that the useful opportunity was not primarily the attention algorithm. Repeated batch-one matrix-vector operations, normalization work, and the cost of launching many small GPU operations were more important. The search then focused on exact CUDA Graph replay around those fixed-size pieces.

Confirmed improvements

The percentages below are relative to the 42.95 tokens/second eager baseline.

1. Capture the MLP decode step in a CUDA Graph

C022 — 48.87 tokens/second (+13.0%)

During one-token decoding, the MLP receives the same shape every time: one batch item, one token, and the model hidden size. The implementation captured that operation once and replayed it, avoiding repeated setup work from Python and the CUDA dispatcher. Prompt processing still used the normal eager path.

2. Capture more of the fixed decode path

C023 — 63.28 tokens/second (+46.3%)

The next version captured the repeated normalization steps and the final output steps as well as the MLP. This removed more small GPU launches from each generated token. Attention and the growing KV cache were deliberately left eager because their shapes and state change during generation.

3. Combine post-attention normalization and the MLP

C024 — 65.91 tokens/second (+52.5%)

Instead of replaying post-attention normalization and the MLP as separate pieces, the implementation captured them together as one block. That removed an intermediate replay and a buffer handoff while preserving the original mathematical operations.

4. Combine input normalization with Q, K, and V projections

C025 — 77.92 tokens/second (+80.1%)

Before attention, the model normalizes the hidden state and then computes the query, key, and value projections. The implementation captured those four operations together for the one-token decode case. The attention calculation, rotary position handling, KV-cache update, and output projection remained unchanged and eager.

The adapter preserved the original Q → K → V call order and checked that order during testing, so the optimization could not silently mix up the three projections.

5. Reuse the prefill normalization

C026 — 78.96 tokens/second (+82.6%)

When reading the initial prompt, the same normalized hidden state was being recomputed before the query, key, and value projections. The final version computed it once and reused it for all three projections, while keeping the C025 decode graphs.

The improvement over C025 itself was small and within timing noise. The important result is that it recovered the prompt-processing path without sacrificing the large exact decode improvement.

Correctness result

The final C026 configuration matched the frozen baseline’s generated token IDs in 30 out of 30 comparison runs across the held-out prompt suite.

The measured progression was:

Eager baseline                 42.95 tok/s
C022 MLP graphs               48.87 tok/s   (+13.0%)
C023 fixed decode graphs      63.28 tok/s   (+46.3%)
C024 compound norm + MLP     65.91 tok/s   (+52.5%)
C025 grouped norm + QKV       77.92 tok/s   (+80.1%)
C026 shared prefill norm     78.96 tok/s   (+82.6%)

What this result does and does not establish

This is an exact-output inference result for MiniCPM5-1B in the tested Transformers/PyTorch environment on an RTX 3060, using batch-one BF16 greedy decoding.

The CUDA Graph technique can potentially be adapted to other NVIDIA GPUs, but graphs must be captured again and each device needs its own baseline and correctness test. The current code is not an Apple Metal, AMD, CPU, vLLM, or SGLang result.

The graph path is written for the one-token decode shape. Other batch sizes and shapes fall back to eager execution unless separate graphs are added. The current QKV coordination was tested as serialized single-request inference, not as a concurrent production server.

The result is therefore best understood as a validated optimization pattern and a working MiniCPM5/RTX 3060 implementation—not as a universal 82.6% guarantee for every model or GPU.

01_core_code_snippets.py

"""Illustrative excerpts from the MiniCPM5 inference optimization.

These snippets show the important ideas, but are not a standalone drop-in
implementation. The actual hook also handles model wiring, graph pools,
warmup streams, capture ordering, and evidence receipts.
"""

import torch


class FixedShapeCudaGraph(torch.nn.Module):
    """Replay a module when decoding exactly one token at batch size one."""

    def __init__(self, eager_module, hidden_size, *, pool=None):
        super().__init__()
        self.eager_module = eager_module
        self.hidden_size = hidden_size

        # CUDA Graph replay needs stable input and output memory addresses.
        self.register_buffer(
            "static_input",
            torch.empty((1, 1, hidden_size), device="cuda", dtype=torch.bfloat16),
            persistent=False,
        )

        # Warm up before capture so lazy CUDA work is finished.
        with torch.inference_mode():
            for _ in range(3):
                eager_module(self.static_input)

        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph, pool=pool), torch.inference_mode():
            static_output = eager_module(self.static_input)
        self.register_buffer("static_output", static_output, persistent=False)

    def forward(self, hidden_states):
        # Other shapes use the original implementation instead of being forced
        # through a graph captured for one token.
        if hidden_states.shape != (1, 1, self.hidden_size):
            return self.eager_module(hidden_states)

        self.static_input.copy_(hidden_states)
        self.graph.replay()
        return self.static_output


class CompoundPostAttentionBlock(torch.nn.Module):
    """Run post-attention normalization and the MLP as one captured block."""

    def __init__(self, eager_norm, eager_mlp, hidden_size, *, pool=None):
        super().__init__()
        self.eager_norm = eager_norm
        self.eager_mlp = eager_mlp
        self.hidden_size = hidden_size
        self.static_input = torch.empty(
            (1, 1, hidden_size), device="cuda", dtype=torch.bfloat16
        )

        def block(x):
            return self.eager_mlp(self.eager_norm(x))

        self.block = block
        with torch.inference_mode():
            for _ in range(3):
                block(self.static_input)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph, pool=pool), torch.inference_mode():
            self.static_output = block(self.static_input)

    def forward(self, x):
        if x.shape != (1, 1, self.hidden_size):
            return self.block(x)
        self.static_input.copy_(x)
        self.graph.replay()
        return self.static_output


def grouped_decode_qkv(x, norm, q_proj, k_proj, v_proj):
    """The grouped decode work captured by C025."""
    normalized = norm(x)
    return q_proj(normalized), k_proj(normalized), v_proj(normalized)


def shared_prefill_qkv(x, norm, q_proj, k_proj, v_proj):
    """The C026 prefill idea: normalize once, reuse three times."""
    normalized = norm(x)
    query = q_proj(normalized)
    key = k_proj(normalized)
    value = v_proj(normalized)
    return query, key, value


def configure_layer(layer, hidden_size, *, pool):
    """Conceptual wiring for one MiniCPM5 decoder layer."""
    # Decode: one graph for input norm + Q/K/V, while attention and KV-cache
    # updates remain eager because their state changes on every token.
    layer.input_layernorm = FixedShapeCudaGraph(
        layer.input_layernorm, hidden_size, pool=pool
    )

    # Decode: one graph for post-attention norm + MLP.
    layer.mlp = CompoundPostAttentionBlock(
        layer.post_attention_layernorm,
        layer.mlp,
        hidden_size,
        pool=pool,
    )


if __name__ == "__main__":
    print("Illustrative snippets only; run the repository harness for evidence.")