perf: vectorize per-sample logit slicing to reduce KL kernel launches (issue #6 item 3) - #20

Merged

marksverdhei merged 3 commits into

marksverdhei:mainfrom

hai-pilgrim:perf/vectorize-logit-loop

Mar 29, 2026

Commits on Mar 29, 2026

perf: vectorize per-sample logit slicing to reduce KL kernel launches (issue #6 item 3)

marksverdhei
and
claude
committed
d2a74ac
refactor: extract _compute_batched_kl helper and add numerical equivalence test

marksverdhei
and
claude
committed
60aa639
style: apply ruff formatting
marksverdhei
committed
6770512