Articles

Python Micro-Optimizations and the AI Infrastructure Stack

8 min read

I started with a tiny question: in Python, is it faster to write arr[i] twice or assign it to a local variable first?

That question is small, but it points at a bigger habit. Performance advice is only useful when you know which layer you are talking about. Python bytecode, C++ compiler output, CUDA kernels, inference engines, and Kubernetes schedulers all have different bottlenecks. The same optimization instinct can be helpful in one layer and noise in another.

Language-Level Performance

At the Python level, source code shape can matter because the interpreter does a lot of work that a compiled language may erase.

Cache Repeated Indexing in Python

If you use the same indexed value more than once, pull it into a local variable:

python
x = arr[i]
f(x)
g(x)

That avoids evaluating arr[i] twice. This is not profound, but in a hot loop it can remove real interpreter work.

If the value is only used once, keep the direct expression:

python
foo(arr[i])

My rule of thumb: do not add a name for one use. Do add one when the same indexed value is used repeatedly.

Runtime Tuples Still Allocate

This expression creates a new tuple object at runtime:

python
(x, y)

CPython makes small allocation relatively cheap through mechanisms like free lists and pymalloc. Tuples are also a little smaller and cheaper than lists. Still, (x, y) is not free. It allocates.

Most of the time that does not matter. It is only worth caring about when tuple creation is inside the part of the program you have already measured.

Two Sum Has a Tight Inner Loop

For the sorted Two Sum problem, I would write the hot loop like this:

python
while True:
    s = numbers[l] + numbers[r]
    if s == target:
        return [l + 1, r + 1]
    if s < target:
        l += 1
    else:
        r -= 1

The useful details are small:

  • Compute numbers[l] + numbers[r] once.
  • Use else instead of elif s > target after the equality check.
  • Remove unreachable fallback code when the problem guarantees a solution.

None of this changes the algorithm. It just keeps repeated interpreter work out of the loop.

else Avoids a Second Python Comparison

In Python:

python
if s < target:
    l += 1
else:
    r -= 1

is slightly cheaper than:

python
if s < target:
    l += 1
elif s > target:
    r -= 1

after you have already checked s == target.

The expected improvement is small: often 0-2%, maybe 2-5% in a very hot loop. The point is not branch prediction. It is avoiding extra bytecode and interpreter dispatch.

Where Python and C++ Diverge

This is where the lesson stops being about Python style. The same source rewrite can mean different things in different languages.

C++ Usually Erases This Kind of Rewrite

In optimized C++, the compiler can often turn this:

cpp
if (x < y) {
  left();
} else if (x > y) {
  right();
}

into the same machine code as this:

cpp
if (x < y) {
  left();
} else {
  right();
}

when the surrounding logic proves equality has already been handled.

That means the typical source-level improvement is 0%. The compiler already knows how to clean up a lot of this.

In C++, the better questions are usually about memory access, cache locality, vectorization, branch predictability, and the algorithm itself.

Python Micro-Optimizations Matter More Than C++ Micro-Rewrites

Python executes source-level operations more literally. Indexing, comparison, attribute lookup, and allocation all go through runtime machinery.

C++ compilers aggressively rewrite, inline, remove redundant work, and optimize across expressions. A cleanup that matters in Python can be irrelevant in C++.

The practical split:

  • In Python, reduce repeated interpreter work in hot loops.
  • In C++, improve the work the CPU actually has to do.

Framework Boundaries Are Not Always Bottlenecks

This also comes up with AI frameworks. The Python/C++ boundary is visible, so it is tempting to blame it. Often it is not where the time is going.

PyTorch vs LibTorch Is Usually Not the Big Lever

PyTorch and LibTorch share the same core backend shape:

text
Python -> ATen -> CUDA kernels
C++    -> ATen -> CUDA kernels

Most model execution time happens in C++ and CUDA either way. Moving from Python PyTorch to C++ LibTorch often gives only small gains:

  • Training: roughly 0-5% in many ordinary cases.
  • Large-model inference: roughly 0-5% when GPU kernels dominate.

The Python overhead is real. It just may not be the bottleneck once execution is inside large GPU kernels.

Inference Engines Are Faster Because of Systems Work

Engines like vLLM, SGLang, TensorRT-LLM, and llama.cpp are not faster just because they are written in C++.

They win through systems-level choices:

  • continuous batching
  • paged attention
  • KV cache management
  • CUDA Graphs
  • FlashAttention
  • custom kernels
  • optimized schedulers

That is a different kind of performance work from shaving a branch in Python.

AI Infrastructure Is a Layered System

Past a single process, performance becomes a coordination problem. The interesting work shifts to batching, memory management, scheduling, quota, storage, and placement.

The Learning Roadmap Moves Down the Stack

The roadmap I would use moves down the stack:

text
Applications
Inference engines
Distributed runtime
GPU architecture
Containers
Kubernetes
Schedulers
Networking
Storage
Operating systems
Hardware

For inference, I would study vLLM, SGLang, and TensorRT-LLM. For training, NCCL, Megatron-LM, DeepSpeed, FSDP, and ZeRO. For GPU execution, CUDA, CUTLASS, FlashAttention, and CUDA Graphs. For storage and systems, checkpointing, object storage, mmap, GPUDirect Storage, RocksDB, Bigtable, Spanner, MapReduce, and Designing Data-Intensive Applications.

Kubernetes Schedules Pods; Kueue Admits Jobs

Normal Kubernetes flow:

text
Job
  -> Pods
  -> Scheduler
  -> Nodes

With Kueue:

text
Job
  -> Kueue
  -> Pods
  -> Scheduler
  -> Nodes

Kueue sits above the Kubernetes scheduler. It decides when a job can start, how quota is enforced, whether unused quota can be borrowed, how priority and preemption work, and whether a large distributed job can be admitted atomically.

The Kubernetes scheduler still decides which specific nodes the admitted pods run on.

The AI Platform Mental Model

A production AI platform is not just PyTorch on GPUs. A more useful mental model looks like this:

text
API
  -> Global load balancer
  -> Kubernetes
  -> Scheduler
  -> Quota system
  -> GPU runtime
  -> CUDA
  -> Inference engine
  -> Model

Each layer has a different job. Request scheduling is not job admission. Job admission is not pod placement. Pod placement is not CUDA execution. CUDA execution is not KV cache management.

The Main Takeaway

The lesson I took from this: first name the layer.

Caching arr[i] is a Python interpreter optimization. Replacing elif with else in C++ is usually a compiler non-event. Moving PyTorch code to LibTorch may not help if CUDA kernels dominate. Faster LLM serving usually comes from batching, memory management, scheduling, and kernel-level execution.

That is the difference between using an ML framework and working on AI infrastructure. You are not just asking "is this faster?" You are asking "which layer is doing the work?"