From 59a36ca76dc1255923ffd2a94a3e3fb13cac0ddd Mon Sep 17 00:00:00 2001 From: seongwoo Date: Tue, 28 Jul 2026 08:13:57 +0900 Subject: [PATCH] [examples] Optimize static Llama last-token logit projection Gather the final valid hidden state before the LM head in prefill and append-prefill paths so the runtime does not materialize full-sequence logits. TICO-DCO-1.0-Signed-off-by: seongwoo --- .../test_static_llama_runtime_helpers.py | 120 +++++++++++++++++- .../recipes/debug/static_llama_runtime.py | 65 +++++++--- 2 files changed, 166 insertions(+), 19 deletions(-) diff --git a/test/quantization/recipes/integration/test_static_llama_runtime_helpers.py b/test/quantization/recipes/integration/test_static_llama_runtime_helpers.py index 701d6891..00105ec6 100644 --- a/test/quantization/recipes/integration/test_static_llama_runtime_helpers.py +++ b/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, :]) diff --git a/tico/quantization/recipes/debug/static_llama_runtime.py b/tico/quantization/recipes/debug/static_llama_runtime.py index f07a6bfb..32698590 100644 --- a/tico/quantization/recipes/debug/static_llama_runtime.py +++ b/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, )