Imported from alpharomercoma/xla-agentic-development (skills/tpu-pallas-debugging/SKILL.md). Install upstream with npx skills add alpharomercoma/xla-agentic-development --skill tpu-pallas-debugging. Copyright stays with the author.

Debugging Pallas TPU kernels

Your role: you are a debugger. Fix the failure, then stop. Do not optimize a kernel that now works, and do not refactor code that was not implicated. A correct slow kernel beats a fast broken one, and a debugging session that quietly becomes a rewrite loses the thread of what actually fixed it.

Triage

Match the failure, go to the section. Do not start editing before you have classified it — the four classes have disjoint fixes and guessing wastes a compile cycle each time.

Reproduce under interpret mode first

Before analyzing anything, try to reproduce on CPU:

pl.pallas_call(kernel, ..., interpret=True)

This single step partitions the problem:

• Reproduces → the bug is logic or indexing. Debug it on CPU, where you have pl.debug_print, real Python tracebacks, and no hardware queue.
• Does not reproduce → the bug is layout, capacity, or precision. Those are exactly the things interpret mode does not model, so its silence is informative.

For synchronization bugs, use the richer simulation, which turns a nondeterministic device hang into a deterministic CPU assertion:

from jax.experimental.pallas import tpu as pltpu

pl.pallas_call(
    kernel,
    interpret=pltpu.InterpretParams(
        detect_races=True,
        out_of_bounds_reads=True,
    ),
    ...
)

Call pltpu.reset_tpu_interpret_mode_state() between runs, or state leaks across tests and you will chase a failure that belongs to the previous case.

Interpret mode has no MXU. It cannot reproduce the bf16 rounding that f32 matmul performs by default. A kernel that is correct under interpret and wrong on device is usually precision, not logic — see Numerics.

Layout errors

The message names a dimension that is not divisible by 8 or 128, or complains about apply-vector-layout.

Only the last two dimensions of a block map onto vector registers, and they must be divisible by 8 (sublanes, for 32-bit) and 128 (lanes), or equal the array's corresponding dimensions. For packed dtypes the sublane requirement is larger: 16 for bf16, 32 for int8.

python3 "${CLAUDE_PLUGIN_ROOT}/skills/tpu-pallas-writing/scripts/check_block_shapes.py" \
  --chip v6e --block 256,512 --dtype bfloat16

Fixes, in order of preference:

1. Round the block up to a legal shape and let ragged edges pad. Cheapest, and usually correct.
2. Reorder the array so the awkward dimension is not last. If a size-3 dimension is trailing, that is a layout problem, not a tiling problem.
3. Reshape outside the kernel. Reshapes that touch the last two dimensions are expensive or unsupported inside it, but free outside where XLA can fuse them.

If a transpose is implicated: transposing over leading axes works when rank ≥ 4; transposing the last two dimensions is the expensive case. Prefer expressing a transposed matmul through dimension_numbers over materializing the transpose.

Capacity errors

RESOURCE_EXHAUSTED at compile time means the working set exceeds VMEM:

sum(block_bytes * buffer_count) + scratch_bytes  <=  VMEM(generation)

buffer_count defaults to 2. VMEM is 16 MiB on v4 and 64–128 MiB on later generations — so a kernel that fits on v5e may not on v5p.

Shrink in this order. Each step costs less performance than the one after it:

1. Halve the last dimension, keeping it a multiple of 128. Cheapest — it reduces the working set without changing the tiling.
2. Halve the second-minor dimension, keeping it a multiple of 8 (or 16/32 for packed dtypes).
3. Drop buffer_count to 1 if you raised it. Costs pipelining overlap.
4. Move a scratch buffer out, recomputing instead of caching.
5. Split the kernel into two pallas_calls.

bisect_block_shape.py automates the search for the largest block that still fits:

python3 "${CLAUDE_PLUGIN_ROOT}/skills/tpu-pallas-debugging/scripts/bisect_block_shape.py" \
  --chip v5p --start 512,1024 --dtype float32 --buffers 2

Synchronization errors

A hang with no useful message, or results that change between runs.

Almost always one of:

• A start() with no matching wait(). Every pltpu.make_async_copy(...).start() needs its wait(). On device this is a timeout; under InterpretParams it is an assertion.
• A remote copy missing wait_send() or wait_recv(). Remote DMA is push-only and both sides must be waited.
• Reading a buffer while a DMA is still writing it. detect_races=True finds this.
• Reusing a semaphore across two overlapping transfers.
• Distinct collectives sharing a collective_id.

Reproduce under InterpretParams(detect_races=True) before touching the code. A hang debugged by inspection is a hang debugged by guessing.

Numerics errors

It compiles, it runs, the numbers are wrong. Classify before fixing:

Wrong under interpret mode too → real logic bug. Indexing, an incorrect index_map, a missing accumulator initialization, a reduction along the wrong axis. Bisect by writing intermediates to separate outputs and comparing stage by stage.

Correct under interpret, wrong on device → almost certainly precision. f32 matmul rounds operands through bf16 by default, and the CPU backend has no MXU, so interpret mode cannot show it. Run the precision sweep:

python3 "${CLAUDE_PLUGIN_ROOT}/skills/tpu-pallas-debugging/scripts/precision_sweep.py" my_kernel.py

Only the first block is right, or only the last → accumulator initialization. The output block is not zeroed on arrival, and scratch persists across grid steps:

@pl.when(pl.program_id(2) == 0)
def _():
    o_ref[...] = jnp.zeros_like(o_ref)

NaN or Inf → usually a missing max-subtraction before exp. See the softmax example in tpu-pallas-writing.

Unsupported operations

Mosaic failed to compile or NotImplementedError often means the operation genuinely is not available. Check tpu-pallas-docs → references/api/unsupported.md before assuming you wrote it wrong.

The usual offenders: int4, integer reductions (cast to f32 and back), reshapes and transposes touching the last two dimensions, and Python-level if on a traced value (use pl.when or jax.lax.cond).

Escalation ladder

Debugging is bounded. Track attempts and escalate on schedule rather than trying variations indefinitely.

Before starting, save a copy: cp kernel.py kernel.py.pre-debug. At attempt 10, stop and write up what was tried, what each attempt changed, what the result was, and what you would try next with more information.

Report the investigations that found nothing. "I ruled out layout, capacity, and synchronization; the failure survives at HIGHEST precision, so it is a logic bug I have not localized" is a genuinely useful result. Silently trying an eleventh variation is not.

Debug flags

jax.config.update("jax_pallas_verbose_errors", True)     # fuller Mosaic diagnostics
jax.config.update("jax_pallas_enable_debug_checks", True) # runtime bounds checks
jax.config.update("jax_pallas_poison_buffers", True)      # surface uninitialized reads

pl.debug_print("i={} v={}", i, value) works inside kernels, including on device.

For PyTorch/XLA callers, PT_XLA_DEBUG_LEVEL=2 plus torch_xla.debug.metrics.metrics_report() shows whether the problem is even in the kernel — see tpu-torch-xla-lowering.

References

• references/interpret-mode.md — InterpretParams in full, and what it cannot model
• references/mosaic-errors.md — error text to root cause
• references/vmem-exhaustion.md — the shrinking ladder with worked numbers
• references/races-and-dma.md — semaphore discipline and deadlock patterns
• references/precision.md — the sweep, and telling policy apart from bugs
• references/SOURCES.md — upstream for every claim

Related skills