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

[hai-pilgrim]

Summary

Addresses issue #6 item 3: "Vectorize per-sample logit slicing loop".

The compute_loss loop called compute_kl_divergence N times per batch (once per sample), each triggering 3 separate CUDA kernel launches (softmax + log_softmax + kl_div). For batch size B=8, that's 24 kernel calls per training step.

This PR reduces it to 3 kernel calls per training step by:

1. Computing per-sample response start positions (CPU, ~microseconds)
2. Building padded batch tensors [valid, max_resp_len, V] via a CPU loop (copies only response slices)
3. Calling compute_kl_divergence once with per_sample=True
4. Averaging per-sample losses

The same vectorization is applied to the sequential_eval path in prediction_step for consistency.

Why it's safe

The aligned batch tensors contain exactly the same data as the per-sample slices — no numerical change to the loss. The per_sample=True KL path was already tested and exists in kl.py.

Test plan

• All existing tests pass (uv run pytest tests/ -k "not benchmark")
• test_compute_loss_returns_scalar — loss is still scalar with gradient
• test_loss_is_differentiable — LoRA params still receive gradients
• test_compute_loss_batch — batch padding logic still correct
• test_compute_loss_empty_inputs / test_compute_loss_empty_responses — zero-loss paths work

🤖 Generated with Claude Code

[marksverdhei]

Review

The vectorization approach is correct — building padded batch tensors and calling compute_kl_divergence once with per_sample=True instead of looping per-sample. Two issues:

1. No new tests for a numerical refactor. This changes the core loss computation path. Please add at least one test that verifies numerical equivalence between old (per-sample loop) and new (batched) approaches on a fixed seed. A regression here would be silent and catastrophic.

2. Duplicated logic between compute_loss and prediction_step. The vectorization pattern (start computation, valid filtering, batch assembly, KL call) is copy-pasted between both methods. Consider extracting it into a shared helper like _compute_batched_kl(teacher_outputs, student_outputs, teacher_inputs, student_inputs, teacher_prompt_lengths, student_prompt_lengths).

Minor: t_seq_len - 1 - (t_starts[i] - 1) simplifies to t_seq_len - t_starts[i].

Please add the equivalence test and consider the dedup. Happy to merge once the test is in.

[marksverdhei]

Reviewed. All requested changes addressed:

1. Numerical equivalence test added — test_batched_kl_matches_per_sample_loop verifies the batched path matches a manual per-sample loop on fixed seed. This is the critical safety net for this refactor.

2. Deduplication done — Extracted _compute_batched_kl, _prepare_pairs, _build_texts_and_lengths, _make_fwd_kwargs as shared helpers. compute_loss and prediction_step now both delegate to these, eliminating the copy-paste.

3. Simplified arithmetic — t_seq_len - t_starts[i] instead of the longer form.

The vectorized approach (building padded batch tensors, single compute_kl_divergence call with per_sample=True) is correct and well-commented. LGTM.