perf: project only sampled rows through lm_head during prefill

- add logits_positions to AutoRegressiveLM.forward, gathering rows before the final norm so the lm_head GEMM covers only the positions prefill samples from
- execute_prefill builds last_token_indices up front and passes them in, dropping the post-forward gather of a [tokens, vocab] tensor
- prefill graph warmup passes a single index; decode stays untouched (every row is sampled) and prefill itself runs eager, so graph capture is unaffected
- update the ragged-prefill fake to slice by the received index and add a packed-row exact-equality test

Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, 512-token prompts, greedy; prefill B=32: 368.1 -> 323.5 ms (44.5k -> 50.6k tok/s, +13.8%), B=8: 89.3 -> 78.9 ms (+13.2%), B=1: 12.3 -> 11.4 ms (+7.9%); decode step unchanged; full suite: 897 passed

astrai/model/transformer.py

@@ -106,6 +106,7 @@ class AutoRegressiveLM(AutoModel):
         kv_cache: Optional[KVCache] = None,
         position_ids: Optional[Tensor] = None,
         fwd: Optional[str] = None,
+        logits_positions: Optional[Tensor] = None,
     ) -> Dict[str, Tensor]:
         if fwd is None:
             if input_ids.ndim != 2:
@@ -142,7 +143,11 @@ class AutoRegressiveLM(AutoModel):
                 aux_losses.append(layer_output["aux_loss"])
                 router_stats_list.append(stats)
 
-        hidden_states = self.norm(x)
+        if logits_positions is not None:
+            # RMSNorm is per-row, so gathering before it matches gathering after.
+            hidden_states = self.norm(x[logits_positions])
+        else:
+            hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)
 
         output = {"logits": logits, "hidden_states": hidden_states}