[examples] Optimize static Llama last-token logit projection

test/quantization/recipes/integration/test_static_llama_runtime_helpers.py

@@ -22,6 +22,44 @@
 HAS_TRANSFORMERS = importlib.util.find_spec("transformers") is not None
 
 
+class _RecordingIdentity(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.input_shape: Optional[tuple[int, ...]] = None
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        self.input_shape = tuple(tensor.shape)
+        return tensor
+
+
+class _PassThroughPrefillLayer(torch.nn.Module):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+    ):
+        del attention_mask, position_embeddings
+        batch_size, seq_len, _ = hidden_states.shape
+        new_k = hidden_states.new_zeros(batch_size, 1, seq_len, 1)
+        new_v = hidden_states.new_zeros(batch_size, 1, seq_len, 1)
+        return hidden_states, new_k, new_v
+
+
+class _FakeReferenceBackbone(torch.nn.Module):
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        *,
+        use_cache: bool,
+        return_dict: bool,
+    ):
+        assert not use_cache
+        assert not return_dict
+        hidden_states = input_ids.to(torch.float32).unsqueeze(-1).repeat(1, 1, 2)
+        return (hidden_states,)
+
+
 class _FixedDeltaAppendLayer(torch.nn.Module):
     def __init__(self, new_k: torch.Tensor, new_v: torch.Tensor):
         super().__init__()
@@ -92,6 +130,82 @@ def test_static_llama_gather_last_token_logits(self):
         self.assertTrue(torch.equal(gathered[0], logits[0, 1]))
         self.assertTrue(torch.equal(gathered[1], logits[1, 3]))
 
+    def test_static_llama_gather_last_valid_token_keeps_sequence_axis(self):
+        """Hidden-state gathering should retain a singleton sequence dimension."""
+        from tico.quantization.recipes.debug.static_llama_runtime import (
+            _gather_last_valid_token,
+        )
+
+        hidden_states = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
+        valid = torch.tensor(
+            [
+                [True, True, False, False],
+                [False, False, True, True],
+            ]
+        )
+
+        gathered = _gather_last_valid_token(hidden_states, valid)
+
+        self.assertEqual(tuple(gathered.shape), (2, 1, 3))
+        torch.testing.assert_close(gathered[0, 0], hidden_states[0, 1])
+        torch.testing.assert_close(gathered[1, 0], hidden_states[1, 3])
+
+    def test_prefill_projects_only_last_valid_hidden_state(self):
+        """Prefill should run final norm and LM head on one sequence position."""
+        from tico.quantization.recipes.debug.static_llama_runtime import (
+            StaticLlamaLayerRuntime,
+        )
+
+        runtime = StaticLlamaLayerRuntime.__new__(StaticLlamaLayerRuntime)
+        runtime.device = torch.device("cpu")
+        runtime.padding_side = "right"
+        runtime.tokenizer = SimpleNamespace(pad_token_id=0)
+        runtime.max_seq = 4
+        runtime.num_hidden_layers = 1
+        runtime.num_kv_heads = 1
+        runtime.head_dim = 1
+        runtime.past_len = 0
+        runtime.embed_tokens = torch.nn.Embedding(32, 2)
+        with torch.no_grad():
+            runtime.embed_tokens.weight.copy_(
+                torch.arange(64, dtype=torch.float32).reshape(32, 2)
+            )
+        runtime.final_norm = _RecordingIdentity()
+        runtime.lm_head = _RecordingIdentity()
+        runtime.rope_cos = torch.zeros(1, runtime.max_seq, 1)
+        runtime.rope_sin = torch.zeros_like(runtime.rope_cos)
+        runtime.layer_caches = []
+        runtime.prefill_layers = torch.nn.ModuleList([_PassThroughPrefillLayer()])
+
+        input_ids = torch.tensor([[5, 6, 0, 0]])
+        attention_mask = torch.tensor([[1, 1, 0, 0]])
+        logits = runtime.prefill(input_ids, attention_mask)
+
+        expected = runtime.embed_tokens(torch.tensor([[6]]))[:, 0, :]
+        self.assertEqual(runtime.past_len, 2)
+        self.assertEqual(runtime.final_norm.input_shape, (1, 4, 2))
+        self.assertEqual(runtime.lm_head.input_shape, (1, 1, 2))
+        self.assertEqual(tuple(logits.shape), (1, 2))
+        torch.testing.assert_close(logits, expected)
+
+    def test_reference_projects_only_last_hidden_state(self):
+        """Reference checks should avoid producing logits for the full sequence."""
+        from tico.quantization.recipes.debug.static_llama_runtime import (
+            StaticLlamaLayerRuntime,
+        )
+
+        runtime = StaticLlamaLayerRuntime.__new__(StaticLlamaLayerRuntime)
+        runtime.model = SimpleNamespace(
+            model=_FakeReferenceBackbone(),
+            lm_head=_RecordingIdentity(),
+        )
+
+        logits = runtime._reference_last_token_logits(torch.tensor([[3, 4, 5]]))
+
+        self.assertEqual(runtime.model.lm_head.input_shape, (1, 1, 2))
+        self.assertEqual(tuple(logits.shape), (1, 2))
+        torch.testing.assert_close(logits, torch.tensor([[5.0, 5.0]]))
+
     def test_append_prefill_attention_mask_layout(self):
         """Append mask should expose valid past and causal tail slots only."""
         from tico.quantization.recipes.debug.static_llama_runtime import (
@@ -131,8 +245,8 @@ def test_append_prefill_writes_only_valid_delta_kv(self):
         runtime.num_hidden_layers = 1
         runtime.past_len = 2
         runtime.embed_tokens = torch.nn.Embedding(32, 2)
-        runtime.final_norm = torch.nn.Identity()
-        runtime.lm_head = torch.nn.Identity()
+        runtime.final_norm = _RecordingIdentity()
+        runtime.lm_head = _RecordingIdentity()
         runtime.rope_cos = torch.zeros(1, runtime.max_seq, 1)
         runtime.rope_sin = torch.zeros_like(runtime.rope_cos)
 
@@ -158,6 +272,8 @@ def test_append_prefill_writes_only_valid_delta_kv(self):
 
         self.assertEqual(runtime.past_len, 4)
         self.assertEqual(tuple(logits.shape), (1, 2))
+        self.assertEqual(runtime.final_norm.input_shape, (1, 3, 2))
+        self.assertEqual(runtime.lm_head.input_shape, (1, 1, 2))
         torch.testing.assert_close(cache.past_k[:, :, :2, :], before_k[:, :, :2, :])
         torch.testing.assert_close(cache.past_v[:, :, :2, :], before_v[:, :, :2, :])
         torch.testing.assert_close(cache.past_k[:, :, 2:4, :], new_k[:, :, :2, :])

tico/quantization/recipes/debug/static_llama_runtime.py

@@ -154,12 +154,30 @@ def _last_valid_token_indices(valid_token_mask: torch.Tensor) -> torch.LongTenso
     return positions.masked_fill(~valid_token_mask, 0).max(dim=1).values
 
 
+def _gather_last_valid_token(
+    tensor: torch.Tensor, valid_token_mask: torch.Tensor
+) -> torch.Tensor:
+    """Gather each batch row's final valid item and keep the sequence dimension."""
+    if tensor.dim() < 2:
+        raise ValueError(
+            f"Expected tensor rank >= 2 with [batch, sequence, ...], got {tensor.dim()}."
+        )
+    if tuple(tensor.shape[:2]) != tuple(valid_token_mask.shape):
+        raise ValueError(
+            "tensor and valid_token_mask must share [batch, sequence] dimensions. "
+            f"Got tensor={tuple(tensor.shape)}, "
+            f"valid_token_mask={tuple(valid_token_mask.shape)}."
+        )
+
+    last_indices = _last_valid_token_indices(valid_token_mask).to(tensor.device)
+    batch_indices = torch.arange(tensor.size(0), device=tensor.device)
+    return tensor[batch_indices, last_indices, ...].unsqueeze(1)
+
+
 def _gather_last_token_logits(
     logits: torch.Tensor, valid_token_mask: torch.Tensor
 ) -> torch.Tensor:
-    last_indices = _last_valid_token_indices(valid_token_mask).to(logits.device)
-    batch_indices = torch.arange(logits.size(0), device=logits.device)
-    return logits[batch_indices, last_indices, :]
+    return _gather_last_valid_token(logits, valid_token_mask).squeeze(1)
 
 
 def _gather_rope_by_position_ids(
@@ -467,8 +485,11 @@ def prefill(
         self.past_len = int(valid.sum(dim=1)[0].item())
 
         hidden_states = self.final_norm(hidden_states)
-        logits = self.lm_head(hidden_states)
-        return _gather_last_token_logits(logits, valid)
+        # LM head is token-wise. Gather first so prefill does not materialize a
+        # [batch, sequence, vocab] logits tensor.
+        last_hidden_state = _gather_last_valid_token(hidden_states, valid)
+        logits = self.lm_head(last_hidden_state)
+        return logits[:, 0, :]
 
     @torch.no_grad()
     def append_prefill(
@@ -596,8 +617,11 @@ def append_prefill(
 
         self.past_len = start_pos + actual_len
         hidden_states = self.final_norm(hidden_states)
-        logits = self.lm_head(hidden_states)
-        return logits[:, actual_len - 1, :]
+        # Compact valid tokens occupy the bucket prefix. Project only the final
+        # valid token instead of producing logits for every bucket position.
+        last_hidden_state = hidden_states[:, actual_len - 1 : actual_len, :]
+        logits = self.lm_head(last_hidden_state)
+        return logits[:, 0, :]
 
     @torch.no_grad()
     def append_prefill_chunked(
@@ -713,6 +737,18 @@ def decode_one(self, input_ids: torch.LongTensor) -> torch.Tensor:
         logits = self.lm_head(hidden_states)
         return logits[:, -1, :]
 
+    @torch.no_grad()
+    def _reference_last_token_logits(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """Compute reference logits without materializing all sequence positions."""
+        base_outputs = self.model.model(
+            input_ids=input_ids,
+            use_cache=False,
+            return_dict=False,
+        )
+        last_hidden_state = base_outputs[0][:, -1:, :]
+        logits = self.model.lm_head(last_hidden_state)
+        return logits[:, 0, :]
+
     @torch.no_grad()
     def generate_greedy(
         self,
@@ -781,8 +817,7 @@ def verify_against_reference(
             device=self.device,
         )
         compact_input_ids = input_ids[valid].reshape(input_ids.size(0), -1)
-        ref_out = self.model(input_ids=compact_input_ids)
-        logits_ref = ref_out.logits[:, -1, :]
+        logits_ref = self._reference_last_token_logits(compact_input_ids)
 
         self._print_diff(
             "Step 0: prefill last-token logits", logits_rt, logits_ref, verbose
@@ -794,8 +829,7 @@ def verify_against_reference(
 
         for step in range(1, steps + 1):
             logits_rt = self.decode_one(next_token)
-            ref_out = self.model(input_ids=generated)
-            logits_ref = ref_out.logits[:, -1, :]
+            logits_ref = self._reference_last_token_logits(generated)
             self._print_diff(
                 f"Step {step}: decode logits", logits_rt, logits_ref, verbose
             )
@@ -874,8 +908,7 @@ def _run_decode_steps_against_reference(
         next_token = torch.argmax(logits, dim=-1, keepdim=True)
         generated = torch.cat([generated, next_token], dim=1)
         logits = runtime.decode_one(next_token)
-        ref_out = runtime.model(input_ids=generated)
-        logits_ref = ref_out.logits[:, -1, :]
+        logits_ref = runtime._reference_last_token_logits(generated)
         runtime._print_diff(
             f"{title_prefix} step {step}",
             logits,
@@ -919,11 +952,10 @@ def _run_multiturn_append_prefill_check(
         device=runtime.device,
     )
     generated = _compact_valid_tokens(first_input_ids, first_valid)
-    ref_out = runtime.model(input_ids=generated)
     runtime._print_diff(
         "Multi-turn step 0: first prefill last-token logits",
         logits,
-        ref_out.logits[:, -1, :],
+        runtime._reference_last_token_logits(generated),
         verbose,
     )
 
@@ -961,11 +993,10 @@ def _run_multiturn_append_prefill_check(
     generated = torch.cat(
         [generated, compact_second_input_ids.to(runtime.device)], dim=1
     )
-    ref_out = runtime.model(input_ids=generated)
     runtime._print_diff(
         f"Multi-turn append-prefill logits (bucket={cfg.append_bucket})",
         logits,
-        ref_out.logits[:, -1, :],
+        runtime._reference_last_token_logits(generated),
         verbose,
     )