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/inference/runtime/executor.py

@@ -131,6 +131,9 @@ def _warmup_cuda_graphs(
                 kv_cache=kv,
                 position_ids=pos_in,
                 fwd="prefill",
+                logits_positions=torch.tensor(
+                    [warmup_len - 1], dtype=torch.long, device=dev
+                ),
             )
         task_cache.task_free(tid)
 
@@ -346,6 +349,11 @@ class Executor:
             ]
         )
 
+        # Last packed position per request; the model projects only these rows.
+        last_token_indices = (
+            torch.tensor(q_lens, dtype=torch.long, device=self.device).cumsum(0) - 1
+        )
+
         with (
             torch.inference_mode(),
             timed(
@@ -363,11 +371,9 @@ class Executor:
                     start_pos=start_pos,
                 ),
                 fwd="prefill",
+                logits_positions=last_token_indices,
             )
-            last_token_indices = (
+            logits = outputs["logits"]
-                torch.tensor(q_lens, dtype=torch.long, device=self.device).cumsum(0) - 1
-            )
-            logits = outputs["logits"][last_token_indices]
 
         step_out, _ = self._sample_logits(logits, tasks, return_logprobs)
         return tasks, step_out

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}

tests/inference/test_scheduler.py

@@ -195,7 +195,11 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
     executor._workspace = MagicMock()
     executor._workspace.max_batch_size = 16  # Add max_batch_size for validation
     all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7)
-    executor.model = MagicMock(return_value={"logits": all_logits})
+
+    def fake_model(ids, *, position_ids, kv_cache, fwd, logits_positions):
+        return {"logits": all_logits[logits_positions]}
+
+    executor.model = MagicMock(side_effect=fake_model)
     executor._sample_logits = MagicMock(
         return_value=([101, 102], torch.tensor([101, 102]))
     )
@@ -210,6 +214,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
     model_args, model_kwargs = executor.model.call_args
     assert model_args[0].tolist() == [11, 12, 21, 22, 23, 24]
     assert model_kwargs["position_ids"].tolist() == [1, 2, 1, 2, 3, 4]
+    assert model_kwargs["logits_positions"].tolist() == [1, 5]
     executor.task_cache.bind.assert_called_once_with(
         ["a", "b"], executor._workspace, start_pos=1
     )

tests/module/test_model_forward.py

@@ -114,6 +114,53 @@ def test_model_forward_contract_uses_dense_training_and_packed_inference():
         )
 
 
+def test_forward_logits_positions_projects_only_requested_rows():
+    """logits_positions gathers packed rows before the lm_head projection."""
+    from astrai.inference.cache import PagePool, TaskCacheManager
+    from astrai.inference.workspace import InferenceWorkspace
+
+    config = AutoRegressiveLMConfig(**TINY_CONFIG)
+    model = AutoRegressiveLM(config).eval()
+    prompts = [[1, 2, 3], [4, 5]]
+    last_rows = torch.tensor([len(prompts[0]) - 1, len(prompts) - 1 + len(prompts[1])])
+
+    pool = PagePool(
+        n_layers=config.num_hidden_layers,
+        n_kv_heads=config.num_key_value_heads,
+        head_dim=config.hidden_size // config.num_attention_heads,
+        max_batch_size=2,
+        max_seq_len=config.max_position_embeddings,
+        device="cpu",
+        dtype=torch.float32,
+    )
+    cache = TaskCacheManager(pool)
+    workspace = InferenceWorkspace(
+        2,
+        config.max_position_embeddings,
+        config.num_attention_heads,
+        config.hidden_size // config.num_attention_heads,
+        torch.device("cpu"),
+        torch.float32,
+    )
+    for tid, ids in zip(("t1", "t2"), prompts):
+        assert cache.task_alloc(tid, ids)
+    input_ids = torch.tensor(sum(prompts, []), dtype=torch.long)
+    position_ids = torch.cat([torch.arange(len(p)) for p in prompts])
+    with torch.inference_mode():
+        kwargs = dict(
+            position_ids=position_ids,
+            kv_cache=cache.bind(["t1", "t2"], workspace, start_pos=0),
+            fwd="prefill",
+        )
+        full = model(input_ids, **kwargs)
+        sliced = model(input_ids, logits_positions=last_rows, **kwargs)
+
+    assert full["logits"].shape == (5, config.vocab_size)
+    assert sliced["logits"].shape == (2, config.vocab_size)
+    assert torch.equal(sliced["logits"], full["logits"][last_rows])
+    assert torch.equal(sliced["hidden_states"], full["hidden_states"][last_rows])
+
+
 def _router_stats(probs, topk_indices):
     return {"probs": probs, "topk_indices": topk_indices}