PyTorch for serving: custom ops and torch.compile

class LlamaMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: QuantizationConfig | None = None,
        bias: bool = False,
        prefix: str = "",
        reduce_results: bool = True,
        disable_tp: bool = False,
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            bias=bias,
            quant_config=quant_config,
            disable_tp=disable_tp,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            reduce_results=reduce_results,
            disable_tp=disable_tp,
            prefix=f"{prefix}.down_proj",
        )
        if hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
        self.act_fn = SiluAndMul()

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
        x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
        x, _ = self.down_proj(x)
        return x

1. Lines 95-101. gate_up_proj is one linear layer for two matrices, the gate and the up projection. output_sizes=[intermediate_size] * 2 stacks them. With tensor parallelism, each rank holds a slice of the columns (Parallelism basics↗).
2. Lines 103-110. down_proj is a row-parallel linear layer. With reduce_results=True, it ends with an all-reduce over the TP ranks.
3. Line 116. SiluAndMul is a CustomOp: a vLLM layer that selects a CUDA kernel or a PyTorch version (vllm/model_executor/layers/activation.py:116).
4. Lines 118-122. forward() calls the three layers. The activation goes through maybe_fused_act_quant(): for a quantized down_proj, it can use one kernel for the activation and the input quantization; else it calls SiluAndMul (vllm/model_executor/layers/fusion/fused_act_quant.py:196-205). The linear layers return a pair; the second value (an optional bias) is not used here.

Shapes. vLLM flattens the batch: the first dimension is the number of tokens in the step, not the number of requests (Batching↗).

Worked example (sizes from the config.json of Llama-3.1-8B: hidden size 4,096, intermediate size 14,336; inferred, not in the repo). A step has 256 tokens and TP = 1, in BF16:

1. x: [256, 4096].
2. gate_up_proj: weight [28672, 4096] (2 × 14,336 rows), output [256, 28672].
3. SiluAndMul: silu(first half) × second half, output [256, 14336].
4. down_proj: weight [4096, 14336], output [256, 4096].

The two weights hold 28,672 × 4,096 + 4,096 × 14,336 = 176,160,768 parameters, which is 352 MB in BF16. With TP = 2, each rank holds half of each weight and computes [256, 14336] after gate_up_proj.

The weights are not loaded in __init__. vLLM makes the modules first and fills the parameters later from the checkpoint; see "Weight loading: from checkpoint files to GPU parameters".

PyTorch has the public decorator torch.library.custom_op. vLLM uses its own helper, because the docstring says that custom_op "can have significant overhead" (vllm/utils/torch_utils.py:1068-1071):

# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
    op_name: str,
    op_func: Callable,
    mutates_args: list[str] | None = None,
    fake_impl: Callable | None = None,
    target_lib: Library | None = None,
    dispatch_key: str | None = None,
    tags: tuple[torch.Tag, ...] = (),
):
    """`torch.library.custom_op` can have significant overhead because it
    needs to consider complicated dispatching logic. This function
    directly registers a custom op and dispatches it to the CUDA backend.
    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
    for more details.

    By default, the custom op is registered to the vLLM library. If you
    want to register it to a different library, you can pass the library
    object to the `target_lib` argument.

    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
    library object. If you want to bind the operator to a different library,
    make sure the library object is alive when the operator is used.
    """
    if mutates_args is None:
        mutates_args = []

    if dispatch_key is None:
        from vllm.platforms import current_platform

        dispatch_key = current_platform.dispatch_key

    schema_str = infer_schema(op_func, mutates_args=mutates_args)

    my_lib = target_lib or vllm_lib
    my_lib.define(op_name + schema_str, tags=tags)
    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
    if fake_impl is not None:
        my_lib._register_fake(op_name, fake_impl)

1. Line 1056. One Library("vllm", "FRAGMENT") object holds all vLLM ops. "FRAGMENT" means that other code can add ops to the same namespace.
2. Lines 1082-1083. mutates_args lists the arguments that the op writes into. The compiler must know them, so it does not reorder or remove the op.
3. Lines 1085-1088. The dispatch key comes from the platform: "CUDA" on NVIDIA (vllm/platforms/cuda.py:230), "CPU" by default (vllm/platforms/interface.py:145).
4. Line 1090. infer_schema reads the type hints of the Python function and writes the op schema, for example (Tensor query, Tensor key, ...) -> ().
5. Lines 1093-1094. define adds the schema; impl binds the Python function for the dispatch key.
6. Lines 1095-1096. A fake implementation only makes empty output tensors with the right shapes. Dynamo calls it during capture instead of the real kernel (general PyTorch behaviour). The helper registers it only when the caller gives one.

Attention as an opaque op. vLLM registers unified_attention_with_output with mutates_args=["output", "output_block_scale"] (vllm/model_executor/layers/attention/attention.py:799-803). It returns nothing; it writes the result into output. Inside, it gets the attention metadata, the layer and the KV cache from the forward context (vllm/model_executor/layers/attention/attention.py:705-723). A second op, unified_kv_cache_update, writes the new keys and values into the cache first. It runs only when the attention backend does not write the KV cache in its own forward (vllm/model_executor/layers/attention/attention.py:564-572). It returns an empty tensor only to make a data dependency: the attention op takes it as an argument, so the compiler keeps the order (vllm/model_executor/layers/attention/attention.py:726-732, vllm/model_executor/layers/attention/attention.py:777-782).

Which platforms use the opaque op. The layer calls torch.ops.vllm.* only when current_platform.opaque_attention_op() is true. Otherwise it calls the Python functions directly (vllm/model_executor/layers/attention/attention.py:436-440, vllm/model_executor/layers/attention/attention.py:542-580). The base class returns false (vllm/platforms/interface.py:1286-1290); CUDA, ROCm, CPU and XPU return true (vllm/platforms/cuda.py:595-596, vllm/platforms/rocm.py:1231-1232, vllm/platforms/cpu.py:640-641, vllm/platforms/xpu.py:609-610).

What Dynamo records. When the compiled function runs for the first time, Dynamo follows its Python bytecode. It records each tensor op as a node of a torch.fx.Graph. The graph has a placeholder node per input, a call_function node per op, and an output node. vLLM's split code walks these nodes (vllm/compilation/backends.py:576-600). Python that does not touch tensors (for example a loop over the layers) runs during capture and leaves only its tensor ops in the graph.

Guards and dynamic shapes. Plain torch.compile adds guards: checks that run before each call and compile again when, for example, a shape changes. vLLM drops the guards in its own modes (vllm/compilation/wrapper.py:126-146). It marks the token dimension as dynamic instead, from dynamic_arg_dims (vllm/model_executor/models/llama.py:337-346). So one graph serves every number of tokens.

A graph break, by example (plain PyTorch, not vLLM code):

import torch

def f(x):
    y = x.sin()
    if y.sum() > 0:        # Python branch on a tensor value
        return y.cos()
    return y

torch.compile(f)(torch.randn(8))                  # works: two graphs and a Python branch
torch.compile(f, fullgraph=True)(torch.randn(8))  # fails: fullgraph does not allow the break

Dynamo cannot know the value of y.sum() during capture. Without fullgraph, it ends the graph before the branch, runs the branch in Python, and starts a new graph. With fullgraph=True, it raises an error instead. vLLM uses fullgraph=True (vllm/compilation/wrapper.py:168-175). So model code that vLLM compiles must not branch on tensor values in forward().

Where the work goes. In the default mode, the backend is VllmBackend (vllm/config/compilation.py:1123-1134). It splits the FX graph at the splitting ops. By default these are the attention ops in _attention_ops (vllm/config/compilation.py:766-788, vllm/config/compilation.py:1184). When Inductor graph partition is off (the O2 default), they also include the two KV cache update ops (vllm/config/compilation.py:1191, vllm/config/compilation.py:1212-1213, vllm/config/vllm.py:311). Then Inductor compiles each piece. The design document says that all compilation finishes before the server takes requests, and that vLLM caches the results (docs/design/torch_compile.md:18-28).

Defaults. The default optimization level is O2 (vllm/config/vllm.py:434). Above O0, an unset mode becomes VLLM_COMPILE (vllm/config/vllm.py:1806-1810). O2 sets the CUDA graph mode FULL_AND_PIECEWISE (vllm/config/vllm.py:310). --enforce-eager sets both the compile mode and the CUDA graph mode to NONE (vllm/config/vllm.py:1745-1747).

CUDA graph modes. NONE, PIECEWISE and FULL are run-time modes. FULL_DECODE_ONLY and FULL_AND_PIECEWISE are pairs: one mode for pure decode batches, one for mixed batches (vllm/config/compilation.py:53-69). Each step, the model runner looks for a captured graph for the padded token count, and falls back to NONE when none fits (vllm/v1/worker/gpu/cudagraph_utils.py:523-553). In the common case (no speculative decoding), a FULL decode graph is captured for batches where each request has 1 new token (vllm/v1/worker/gpu/model_runner.py:442-445, vllm/v1/worker/gpu/cudagraph_utils.py:333-339). A batch matches it only with that uniform length and at most the captured token count (vllm/v1/worker/gpu/cudagraph_utils.py:125-138). Other batches use PIECEWISE graphs.

Worked example (illustrative numbers). A decode step of a 32-layer model starts about 15 kernels per layer in eager mode: 480 launches. At about 5 µs of CPU time per launch, that is 2.4 ms per step. A step of a small model on a fast GPU can take less than that, so the GPU would wait for the CPU. With a FULL CUDA graph, the CPU starts one replay instead.