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.

API Conventions

Read API patterns before adding or changing the public wrapper, backend selection, block-size config, or input normalization contract. Keep the reference/XLA path usable even when accelerator-specific constraints are not met. Keep backend-specific validation in backend-specific modules.

Cost Estimate Requirement

Add cost_estimate= to each pl.pallas_call:

• Use pl.estimate_cost on a body-equivalent JAX function, not a kernel body with pl.program_id.
• Include IO bytes from call inputs/outputs.

from levanter.kernels.pallas.cost_estimate_utils import with_io_bytes_accessed


def _cost_estimate(
    q: jax.Array,
    k: jax.Array,
    v: jax.Array,
    *,
    kernel_inputs_specs,
    kernel_outputs_specs,
) -> pl.CostEstimate | None:
    body_cost = pl.estimate_cost(reference_impl, q, k, v)
    return with_io_bytes_accessed(
        body_cost,
        kernel_inputs_specs=kernel_inputs_specs,
        kernel_outputs_specs=kernel_outputs_specs,
    )

Definition of Done

• Values match reference within tolerance on the tested grid.
• Gradients match reference on small shapes.
• CPU/reference and accelerator fast paths are covered by tests where applicable.
• Public API, fallback semantics, block-size config, and tuned table behavior match API patterns.
• Every Pallas, Mosaic, or FFI kernel call is inside an explicit shard_map, or its wrapper documents and tests why the inputs are fully local or replicated. Tests or profile evidence show it did not lower through unintended all-gathers.
• Each pl.pallas_call has a reviewed cost_estimate=.
• Benchmark/tuning artifacts include the required schema from Performance workflow.
• Roofline performance is within expected bounds, or limitations are explicitly documented.
• Performance improves on at least one realistic target shape, or limitations are explicitly documented.
• Tuned table is checked in for requested hardware/shape regimes.
• Long-running research records follow the research workflow.