Computational Graphs and Operators

1. What is it?#

A model is a directed acyclic graph: nodes are operators (matmul, add, softmax), edges are tensors. Inference is executing that graph in topological order.

   input
     │
     ▼
  [MatMul]◄── W1
     │
     ▼
   [Add]◄──── b1
     │
     ▼
   [SiLU]
     │
     ▼
  [MatMul]◄── W2
     │
     ▼
   output

An operator (op) is a named computation with defined semantics — MatMul, Softmax, LayerNorm. A kernel is a concrete implementation of an operator for a specific device, dtype, and shape regime. One operator, many kernels.

2. Why does it exist?#

Because the graph is a program you can analyze before running it. Since there is no data- dependent control flow (Section I.01), you can:

• Plan memory allocation exactly.
• Fuse adjacent operators.
• Reorder and eliminate operations.
• Pick optimal kernels per shape.
• Compile the whole thing ahead of time.

None of that is possible for a general program. It is possible here, and that possibility is the basis of TensorRT, TorchInductor, ONNX Runtime, XLA, and TVM.

3. Simple analogy#

A recipe as a flowchart. “Chop onions” and “heat oil” are independent — you can do them in either order or simultaneously. “Fry onions” depends on both. A cook who sees the dependency graph rather than a linear list can parallelize, batch steps that use the same pan, and avoid washing up between compatible operations. That’s exactly what a graph compiler does.

4. Tiny example#

// graph.go — a model is a graph of operators. Here is one, built and printed by hand.
package main

import (
	"fmt"
	"strings"
)

type Node struct {
	Name   string
	Op     string // "input", "matmul", "add", "silu", ...
	Inputs []*Node
}

type Graph struct{ Nodes []*Node }

func (g *Graph) Add(name, op string, inputs ...*Node) *Node {
	n := &Node{name, op, inputs}
	g.Nodes = append(g.Nodes, n)
	return n
}

func (g *Graph) String() string {
	var b strings.Builder
	for _, n := range g.Nodes {
		var in []string
		for _, i := range n.Inputs {
			in = append(in, "%"+i.Name)
		}
		fmt.Fprintf(&b, "%%%-4s = %s(%s)\n", n.Name, n.Op, strings.Join(in, ", "))
	}
	return b.String()
}

func main() {
	// h = silu(x·w1 + b1);  out = h·w2
	g := &Graph{}
	x, w1, b1, w2 := g.Add("x", "input"), g.Add("w1", "input"), g.Add("b1", "input"), g.Add("w2", "input")
	h := g.Add("h0", "matmul", x, w1)
	h = g.Add("h1", "add", h, b1)
	h = g.Add("h2", "silu", h)
	g.Add("out", "matmul", h, w2)

	fmt.Print(g) // see the graph
}

Output:

%x    = input()
%w1   = input()
%b1   = input()
%w2   = input()
%h0   = matmul(%x, %w1)
%h1   = add(%h0, %b1)
%h2   = silu(%h1)
%out  = matmul(%h2, %w2)

Frameworks build exactly this structure for you by tracing your model code (PyTorch’s torch.fx.symbolic_trace prints a graph that looks almost identical).

Four operators. Now note: matmul → add → silu are three separate kernels, three round trips to HBM for the intermediate. A compiler can fuse add and silu into the matmul’s epilogue, turning 3 kernels into 1 (file 09).

5. Technical explanation#

The operator set of an LLM#

Surprisingly small. A complete Llama-family model uses:

MatMul / Linear          the bulk of FLOPs
Embedding (Gather)       lookup
RMSNorm                  reduction + elementwise
SiLU / GELU              elementwise
Mul / Add                elementwise (residual, gating)
RoPE (Rotate)            elementwise with sin/cos tables
Attention                the composite: matmul, scale, mask, softmax, matmul
Softmax                  reduction + elementwise
TopK / Sort / Multinomial  sampling
Cast                     dtype conversion
Reshape / Transpose      metadata only
Concat / Slice           KV cache management

About 12 distinct operators. That’s why writing a working inference engine from scratch (Project 04) is achievable, and why hardware vendors can support LLMs with a modest kernel library.

Graph representations you will meet#

Static vs dynamic graphs#

Static (define-then-run):    build the whole graph, then execute repeatedly
    ✓ full optimization, memory planning, AOT compilation
    ✗ shapes must be known or bucketed; control flow is awkward

Dynamic (define-by-run):     execute ops as Python runs them
    ✓ easy to write and debug, arbitrary control flow
    ✗ per-op dispatch overhead, no cross-op optimization

PyTorch is dynamic by default; torch.compile captures a static graph opportunistically and falls back to eager on “graph breaks” (file 10).

What the graph enables#

1. Memory planning. Knowing the whole graph, you can compute each tensor’s live range and reuse buffers:

Naive:   allocate every intermediate         → peak = sum of all
Planned: reuse buffers whose live ranges don't overlap → peak = max concurrent

For a transformer this can halve activation memory.

2. Constant folding. scale = 1/sqrt(d_head) computed once at build time, not per call.

3. Dead code elimination. Dropout in eval mode, unused outputs, debugging branches.

4. Layout assignment. Choose the tensor layout that minimizes total transpose cost across the whole graph, not greedily per op.

5. Kernel selection. For each matmul, pick the best cuBLAS algorithm for that exact shape, measured at build time (TensorRT literally benchmarks candidates).

6. Under the hood#

How PyTorch eager dispatches one op:

python: x @ w
  → __matmul__ → torch.matmul
  → C++ dispatcher: look up (op, dtype, device, layout) in the dispatch table
  → autograd layer (skipped under inference_mode)
  → backend kernel: at::native::matmul_cuda
  → cuBLAS call
  → cudaLaunchKernel

Roughly 2-10 µs of CPU per op, before any GPU work. With ~350 ops per token, that’s the launch overhead from file 01.

A compiled graph collapses this: one Python call, a pre-built sequence of kernel launches (or a single CUDA graph replay).

7. Performance implications#

• Eager mode costs 2-10 µs per op. At small batch this dominates.
• A captured graph enables CUDA graphs, removing nearly all of it.
• Fusion opportunities are only visible at the graph level. An op-at-a-time executor cannot fuse.
• Memory planning at the graph level reduces peak activation memory substantially.

8. Production implications#

• Use a compiled path in production (torch.compile, TensorRT-LLM, or your engine’s built-in graph capture). Eager mode is for development.
• Compilation is not free: minutes of build time, and it must be cached and keyed by (model, GPU arch, shapes, library versions). Do it at image build time or on first start with a persistent cache (Section II.07).
• Graph breaks kill the benefit. A single unsupported Python construct in the hot path splits the graph and reintroduces overhead. Check with TORCH_LOGS="graph_breaks" python ....
• Dynamic shapes complicate everything (file 11) — this is the central tension in LLM compilation.

9. Common mistakes#

Assuming eager and compiled produce identical numbers. Fusion changes operation order and therefore rounding. Differences are tiny but nonzero, and after sampling can change the output text.

Compiling at startup in production. Adds minutes to cold start.

Ignoring graph breaks. A print() or a .item() in the hot loop and you’ve lost the compilation benefit.

Building a TensorRT engine on one GPU and deploying to another architecture. Engines are architecture-specific.

10. Hands-on exercise#

A. Trace a model. Use torch.fx.symbolic_trace on a small transformer. Print the graph. Count operators by type. Which are matmuls, which are elementwise?

B. Count dispatch overhead. Time a forward pass in eager mode with torch.profiler, separating CPU time from GPU time. What fraction is dispatch?

C. Compile it. Run the same model with torch.compile. Measure the speedup at batch 1 and batch 64. Explain why they differ.

D. Find graph breaks. Add a print(x.sum().item()) inside the model’s forward. Run with TORCH_LOGS="graph_breaks" and observe. Measure the performance cost.

E. Export to ONNX. Export a small model and visualize with Netron. Compare the operator list to the one in section 5.

11. Interview questions#

1. What is a computational graph and why does having one enable optimization?
2. Difference between an operator and a kernel?
3. Name four optimizations that require graph-level visibility.
4. What is a graph break in torch.compile and why does it matter?
5. Why is eager mode slower at batch 1 than at batch 256, proportionally?
6. What must a TensorRT engine be keyed by, and why?

12. Further reading#

• [REFERENCE] PyTorch torch.fx and torch.compile documentation
• [ESTABLISHED] Chen et al., “TVM” (OSDI 2018)
• [REFERENCE] ONNX operator specification
• Next: 03 — GEMM and GEMV