Add Pallas Kernel

Add or update a Pallas kernel

How to apply this skill

Load only the detail files needed for the requested work:

• Kernel sources: read when choosing an in-repo or external kernel to imitate.
• Performance workflow: read before benchmarking, profiling, roofline analysis, or autotuning.
• API patterns: read before adding or changing a public kernel wrapper, fallback order, or block-size config.
• TPU tips: read for TPU Pallas/Mosaic kernels, TPU-specific lowering failures, scoped VMEM, or TPU compiler dumps.
• GPU tips: read for GPU Pallas/Mosaic work.
• Deep references live under docs/reference/; read them only when the routed detail files point there.

Use research only when the user explicitly requests a multi-session research workflow.

Kernel Deliverables

For a kernel K, produce:

• Vanilla JAX reference and Pallas wrapper with the same public API.
• Value, gradient, CPU, and applicable accelerator parity harness.
• Explicit backend and shape validation, with tests for ordered implementation selection and each fallback path.
• Roofline estimate and steady-state benchmark on representative shapes/dtypes.
• When tuning is requested, bounded autotuning, a checked-in tuned table, explicit fallback, and cached autotune-on-miss results.

Correctness Workflow

1. Start from a reference

Use an existing in-repo implementation, pseudocode, a PyTorch reference, or a JAX baseline. The baseline must be obvious and stable, not clever. If the naive baseline would materialize huge intermediates, use a streaming/blockwise baseline with identical math.

2. Write a value and gradient harness

Minimum checks:

• Value parity over a shape/dtype grid.
• Gradient parity on small shapes.
• Backend numerics on CPU and accelerator backends as applicable.
• Pointwise deviation metrics such as max/mean absolute diff, not only allclose.

Use explicit shape/dtype annotations for public APIs and references, such as jaxtyping, where available.

3. Promote long-lived checks to pytest

For in-tree kernels, add or extend tests under lib/levanter/tests/kernels/. Compare the default implementation against the reference on small CPU shapes and accelerator-aligned shapes for fast paths. Read TESTING.md and the nearest module AGENTS.md before writing or changing tests.

Pallas Kernel Workflow

Once the reference is correct, design the Pallas implementation. Use the reference as both a correctness oracle and a performance baseline.

Use existing kernels for structure and API inspiration. Read Kernel sources unless the user already named the specific kernel to follow. Unless there is a stronger local pattern, start by reimplementing the reference in Pallas.

Wrap accelerator kernel boundaries in an explicit jax.shard_map by default. This applies to pl.pallas_call, Mosaic GPU kernels, and custom FFI calls. Reshard inputs to the intended local PartitionSpec before the shard_map, keep the sequence or other nonlocal dimensions unsharded unless the kernel is explicitly written for them, and add a regression check that the lowered JAXPR or HLO contains the expected shard_map. Do not rely on XLA to infer a good sharding for an opaque kernel call boundary. Exceptions are limited to wrappers whose inputs are explicitly documented and tested as fully local or replicated.

Check correctness against the harness and reference implementation before tuning. Once the kernel is correct, run a performance harness on representative shapes/dtypes and compare against the roofline. If performance is not near the expected roofline, read Performance workflow and investigate compiler dumps, pressure signals, and tile choices before broad rewrites.