JAX Pathways

Course · Pallas kernelsPhase 13 · Study guide

Does a custom kernel earn its complexity?

Systems Specialization · CPU ReferenceCPU reference; target qualification separate4 lessons (~7.4h)

Revision v0.7.0 ·

A lower-level implementation first owes you the same numerical contract. Use hostile shapes, independent references and explicit target limits before attempting a speed claim.

Preparation: Distributed training

Runtime: CPU interpretation/simulation reference; TPU target lab requires separately validated hardware

Start the phaseSee the project

Check your starting point

A matrix dimension is not divisible by the tile size. What must the last tile do?

Compare your reasoning

Mask or pad out-of-range loads and stores under a declared policy. Valid elements still need exactly one correct output; fewer launched programs do not excuse missing boundary elements.

If this is unfamiliar, work through Pallas grids and BlockSpecs before attempting the project.

Build in stages

1. Prove indexing and arithmetic

Map each grid coordinate to input/output slices, test rectangular and boundary shapes and reject unsupported layouts or dtypes before launch.

2. Measure only a supported target

Check pipeline semantics and buffer assumptions, compare with a strong compiled baseline and make an end-to-end decision from the operation’s measured share.

Try a changed condition

Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.

Interpret mode passes all numerical tests. Can you use its timings to choose the fastest TPU tile?

Compare an approach

No. Interpretation supports bounded semantic checks. Compile and measure supported candidates on the named TPU with equivalent work and synchronization, including numerical verification there.

Diagnose a likely failure

Symptom: A kernel passes square inputs but corrupts rectangular ones.

Check next: Inspect operand-specific index maps and boundary masks; square dimensions can hide an axis swap.

Decide what is ready

Use kernel-audit. Retain independent numerical checks, rejected unsupported contracts and a target measurement plan. Keep hardware timings unmeasured until actually executed.

Where to put your project code

Write your code in: projects/kernel-audit/my_model.py (copy projects/kernel-audit/starter/model.py first: cp projects/kernel-audit/starter/model.py projects/kernel-audit/my_model.py).

Verify your implementation: python3 projects/kernel-audit/tests/check.py --stage all --impl projects/kernel-audit/my_model.py

Run the staged project checks on your implementation in projects/kernel-audit/my_model.py and record your transfer experiment results in your portfolio.

Open A checked Pallas kernel with an explicit target qualification plan

What still needs further work

Actual TPU/GPU compilation, race diagnosis and target tuning remain unqualified; CPU interpretation cannot close these gaps.