[levanter] Share Pallas autotune helpers and restore compile offload

[dlwh]

Move shard-aware autotune benchmarking out of fused cross-entropy into a shared Pallas helper and restore compile offload for shard-mapped autotune sweeps. This keeps fused CE behavior intact while making the benchmark path reusable from other kernels.

Fixes #4129

[dlwh]

@claude review

[claude]

Claude Code is working…

I'll analyze this and get back to you.

View job run

[dlwh]

@dlwh-golem will you review please

[dlwh-golem]

No findings.



I reviewed the helper extraction in lib/levanter/src/levanter/kernels/pallas/autotune_utils.py and the fused CE integration in lib/levanter/src/levanter/kernels/pallas/fused_cross_entropy_loss/api.py. The extraction preserved the prior shard-map/manual-sharding/tracer behavior, and the new generic tests cover the moved logic plus the fused CE autotune sweep path.



I also ran:

uv run --project lib/levanter --group test python -m pytest -q lib/levanter/tests/kernels/test_pallas_autotune_utils.py lib/levanter/tests/kernels/test_pallas_fused_cross_entropy_loss.py



Result: 68 passed, 12 skipped in about 9s. Residual risk is low; the only thing not really proven yet is behavior for future non-CE adopters with more complex out_specs, but nothing in this PR looks blocking.

Codex work

• Repository: marin-community/marin
• Trigger: [levanter] Share Pallas autotune helpers and restore compile offload #4130 (comment)
• Request ID: 45aad31d-3bc8-4e6a-834c-25a79966b7cb
• Accepted: true

Codex transcript

Completed items:
- none captured

Transcript:
No findings.

I reviewed the helper extraction in [lib/levanter/src/levanter/kernels/pallas/autotune_utils.py](/home/golem/oa-jder-bot/worktrees/github.com/marin-community/marin/4130/lib/levanter/src/levanter/kernels/pallas/autotune_utils.py) and the fused CE integration in [lib/levanter/src/levanter/kernels/pallas/fused_cross_entropy_loss/api.py](/home/golem/oa-jder-bot/worktrees/github.com/marin-community/marin/4130/lib/levanter/src/levanter/kernels/pallas/fused_cross_entropy_loss/api.py). The extraction preserved the prior shard-map/manual-sharding/tracer behavior, and the new generic tests cover the moved logic plus the fused CE autotune sweep path.

I also ran:
`uv run --project lib/levanter --group test python -m pytest -q lib/levanter/tests/kernels/test_pallas_autotune_utils.py lib/levanter/tests/kernels/test_pallas_fused_cross_entropy_loss.py`

Result: `68 passed, 12 skipped` in about 9s. Residual risk is low; the only thing not really proven yet is behavior for future non-CE adopters with more complex `out_specs`, but nothing in this PR looks blocking.

[yonromai]

Approved. I did not find a clear regression in the helper extraction or in the restored compile-offload path. The shared helper preserves the prior fused CE behavior around manual sharding, shard-map wrapping, tracer lowering, and fallback handling, and the targeted tests cover both the new helper and the fused CE sweep path.

Generated with Codex