marin: Skill for Claude Code

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-engineering for repeated tuning launches or controlled comparisons expected to produce reusable evidence. Ordinary implementation and correctness debugging do not need a research notebook.

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.