| |
| |
|
|
| """Reusable teacher-forcing and benchmark run helpers for demos and tests.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import time |
| from dataclasses import dataclass, field |
| from pathlib import Path |
|
|
| import torch |
| from loguru import logger |
|
|
| import ttnn |
|
|
| _SAME_SAMPLING_PARAMS = object() |
|
|
|
|
| def make_contiguous_page_table(batch_size: int, max_seq_len: int, block_size: int = 32) -> torch.Tensor: |
| """Create a contiguous demo page table with one disjoint block range per user.""" |
| if min(batch_size, max_seq_len, block_size) <= 0: |
| raise ValueError("page-table dimensions must be positive") |
| blocks_per_user = (max_seq_len + block_size - 1) // block_size |
| return torch.arange(batch_size * blocks_per_user, dtype=torch.int32).reshape(batch_size, blocks_per_user) |
|
|
|
|
| @dataclass |
| class TeacherForceResult: |
| """Result from a teacher-forcing evaluation run.""" |
|
|
| predicted_tokens: list[int] |
| predicted_tokens_per_user: list[list[int]] |
| reference_top5: torch.Tensor |
| prefill_time_s: float = 0.0 |
| compile_decode_time_s: float = 0.0 |
| decode_times_s: list[float] = field(default_factory=list) |
| batch_size: int = 1 |
| prefill_len: int = 0 |
|
|
| def top1_accuracy(self) -> float: |
| matches = sum( |
| 1 for i, prediction in enumerate(self.predicted_tokens) if self.reference_top5[i, 0].item() == prediction |
| ) |
| return matches / len(self.predicted_tokens) |
|
|
| def top5_accuracy(self) -> float: |
| matches = sum( |
| 1 for i, prediction in enumerate(self.predicted_tokens) if prediction in self.reference_top5[i, :] |
| ) |
| return matches / len(self.predicted_tokens) |
|
|
| @property |
| def ttft_ms(self) -> float: |
| return self.prefill_time_s / self.batch_size * 1000 if self.batch_size else 0.0 |
|
|
| @property |
| def prefill_time_to_token_s(self) -> float: |
| return self.prefill_time_s / self.batch_size if self.batch_size else 0.0 |
|
|
| @property |
| def prefill_tok_s(self) -> float: |
| return (self.batch_size * self.prefill_len) / self.prefill_time_s if self.prefill_time_s > 0 else 0.0 |
|
|
| @property |
| def decode_tok_s_u(self) -> float: |
| return len(self.decode_times_s) / sum(self.decode_times_s) if self.decode_times_s else 0.0 |
|
|
| @property |
| def decode_tok_s(self) -> float: |
| return self.decode_tok_s_u * self.batch_size |
|
|
|
|
| @dataclass |
| class PerfBenchmarkResult: |
| """Result from a performance benchmark run.""" |
|
|
| prefill_time_s: float |
| compile_decode_time_s: float |
| decode_times_s: list[float] |
| batch_size: int |
| num_decode_tokens: int |
| generated_token_ids: list[list[int]] |
| decode_iteration_times_s: list[float] = field(default_factory=list) |
| argmax_top2_margins: list[list[float]] | None = None |
|
|
| @property |
| def ttft_ms(self) -> float: |
| """TTTv1-style average time to first token per user.""" |
| return self.prefill_time_s / self.batch_size * 1000 |
|
|
| @property |
| def tok_s_u(self) -> float: |
| """Tokens per second per user during steady-state decode.""" |
| if not self.decode_times_s: |
| return 0.0 |
| return len(self.decode_times_s) / sum(self.decode_times_s) |
|
|
| @property |
| def tok_s(self) -> float: |
| """Total decode throughput.""" |
| return self.tok_s_u * self.batch_size |
|
|
| @property |
| def decode_latency_mean_ms(self) -> float: |
| if not self.decode_times_s: |
| return 0.0 |
| return sum(self.decode_times_s) / len(self.decode_times_s) * 1000 |
|
|
| def meets_target(self, expected: dict, tolerance: float = 0.05) -> dict[str, bool]: |
| """Check benchmark metrics against the expected thresholds.""" |
| return { |
| "tok_s_u": self.tok_s_u >= expected["tok_s_u"] * (1 - tolerance), |
| "ttft_ms": self.ttft_ms <= expected["ttft_ms"] * (1 + tolerance), |
| } |
|
|
|
|
| def _compile_prefill_and_decode( |
| execution_target, |
| *, |
| prefill_tokens: torch.Tensor, |
| prefill_page_table: torch.Tensor, |
| kv_cache=None, |
| prompt_lens: torch.Tensor | None = None, |
| empty_slots: list[int] | None = None, |
| start_pos: torch.Tensor | None = None, |
| sampling_params=None, |
| prefill_sampling_params=_SAME_SAMPLING_PARAMS, |
| decode_tokens: torch.Tensor | None = None, |
| decode_start_pos: torch.Tensor | None = None, |
| decode_page_table: torch.Tensor | None = None, |
| ) -> None: |
| """Compile the concrete prefill and decode cases through the public target surface.""" |
| assert prefill_tokens.dim() == 2, f"prefill_tokens must be [batch_size, seq_len], got {prefill_tokens.dim()}D" |
| assert ( |
| prefill_page_table.dim() == 2 |
| ), f"prefill_page_table must be [batch_size, max_blocks], got {prefill_page_table.dim()}D" |
|
|
| batch_size = prefill_tokens.shape[0] |
| if decode_tokens is None: |
| decode_tokens = torch.zeros(batch_size, dtype=torch.long, device=prefill_tokens.device) |
| if decode_start_pos is None: |
| decode_start_pos = torch.full( |
| (decode_tokens.shape[0],), |
| prefill_tokens.shape[-1], |
| dtype=torch.long, |
| device=prefill_tokens.device, |
| ) |
| if decode_page_table is None: |
| decode_page_table = prefill_page_table |
|
|
| if prefill_sampling_params is _SAME_SAMPLING_PARAMS: |
| prefill_sampling_params = sampling_params |
|
|
| if sampling_params is not None: |
| execution_target.compile_decode( |
| tokens=decode_tokens, |
| start_pos=decode_start_pos, |
| page_table=decode_page_table, |
| kv_cache=kv_cache, |
| sampling_params=sampling_params, |
| ) |
| execution_target.compile_prefill( |
| tokens=prefill_tokens, |
| page_table=prefill_page_table, |
| kv_cache=kv_cache, |
| prompt_lens=prompt_lens, |
| empty_slots=empty_slots, |
| start_pos=start_pos, |
| sampling_params=prefill_sampling_params, |
| ) |
| return |
|
|
| execution_target.compile_prefill( |
| tokens=prefill_tokens, |
| page_table=prefill_page_table, |
| kv_cache=kv_cache, |
| prompt_lens=prompt_lens, |
| empty_slots=empty_slots, |
| start_pos=start_pos, |
| sampling_params=None, |
| ) |
| execution_target.compile_decode( |
| tokens=decode_tokens, |
| start_pos=decode_start_pos, |
| page_table=decode_page_table, |
| kv_cache=kv_cache, |
| sampling_params=None, |
| ) |
|
|
|
|
| def _profiler_start(profiler, name: str) -> None: |
| if profiler is not None: |
| profiler.start(name) |
|
|
|
|
| def _profiler_end(profiler, name: str) -> None: |
| if profiler is not None: |
| profiler.end(name) |
|
|
|
|
| def run_teacher_forcing( |
| executor, |
| *, |
| prompt_tokens: torch.Tensor, |
| reference_tokens: torch.Tensor, |
| top5_tokens: torch.Tensor, |
| kv_cache: list, |
| page_table: torch.Tensor, |
| max_batch_size: int = 1, |
| profiler=None, |
| ) -> TeacherForceResult: |
| """Run teacher-forcing accuracy measurement against an execution target.""" |
| execution_target = executor |
| batch_size = prompt_tokens.shape[0] |
| assert ( |
| batch_size == max_batch_size |
| ), f"Teacher forcing expects active batch to match max_batch_size, got {batch_size} vs {max_batch_size}" |
| prompt_len = prompt_tokens.shape[-1] |
| num_target = len(reference_tokens) - prompt_len |
| prompt_lens = torch.tensor([prompt_len] * batch_size) |
| empty_slots = list(range(batch_size)) |
|
|
| _compile_prefill_and_decode( |
| execution_target, |
| prefill_tokens=prompt_tokens, |
| prefill_page_table=page_table, |
| kv_cache=kv_cache, |
| prompt_lens=prompt_lens, |
| empty_slots=empty_slots, |
| ) |
|
|
| logger.info(f"Teacher forcing: prefilling {prompt_len} tokens with batch={batch_size}") |
| _profiler_start(profiler, "inference_prefill") |
| try: |
| start_time = time.perf_counter() |
| prefill_output = execution_target.prefill_forward( |
| prompt_tokens, |
| page_table=page_table, |
| kv_cache=kv_cache, |
| prompt_lens=prompt_lens, |
| empty_slots=empty_slots, |
| ) |
| _synchronize_target(execution_target) |
| prefill_time_s = time.perf_counter() - start_time |
| finally: |
| _profiler_end(profiler, "inference_prefill") |
| first_tokens = torch.argmax(prefill_output, dim=-1).view(-1).tolist() |
| predicted_tokens_per_user = [[int(token)] for token in first_tokens] |
|
|
| logger.info(f"Teacher forcing: decoding {num_target - 1} tokens") |
| compile_decode_time_s = 0.0 |
| decode_times_s = [] |
| _profiler_start(profiler, "inference_decode") |
| try: |
| for step in range(1, num_target): |
| ground_truth_token = reference_tokens[prompt_len + step - 1] |
| decode_token = torch.full((batch_size,), ground_truth_token, dtype=torch.long) |
| current_pos = torch.full((batch_size,), prompt_len + step - 1, dtype=torch.long) |
| start_time = time.perf_counter() |
| logits, _ = execution_target.decode_forward( |
| decode_token, |
| current_pos, |
| page_table=page_table, |
| kv_cache=kv_cache, |
| read_from_device=True, |
| ) |
| elapsed = time.perf_counter() - start_time |
| if step == 1: |
| compile_decode_time_s = elapsed |
| else: |
| decode_times_s.append(elapsed) |
|
|
| next_tokens = torch.argmax(logits[:, -1, :], dim=-1).view(-1).tolist() |
| for user_id, token in enumerate(next_tokens): |
| predicted_tokens_per_user[user_id].append(int(token)) |
| finally: |
| _profiler_end(profiler, "inference_decode") |
|
|
| return TeacherForceResult( |
| predicted_tokens=predicted_tokens_per_user[0], |
| predicted_tokens_per_user=predicted_tokens_per_user, |
| reference_top5=top5_tokens[:num_target], |
| prefill_time_s=prefill_time_s, |
| compile_decode_time_s=compile_decode_time_s, |
| decode_times_s=decode_times_s, |
| batch_size=batch_size, |
| prefill_len=prompt_len, |
| ) |
|
|
|
|
| def _split_output(output): |
| return output if isinstance(output, tuple) else (output, None) |
|
|
|
|
| def _target_mesh_device(execution_target): |
| return getattr(execution_target, "mesh_device", None) |
|
|
|
|
| def _target_cluster_shape(execution_target): |
| cluster_shape = getattr(execution_target, "cluster_shape", None) |
| if cluster_shape is not None: |
| return list(cluster_shape) |
| mesh_device = _target_mesh_device(execution_target) |
| return list(mesh_device.shape) if mesh_device is not None else [1, 1] |
|
|
|
|
| def _synchronize_target(execution_target): |
| mesh_devices = getattr(execution_target, "mesh_devices", None) |
| if mesh_devices is not None: |
| for mesh_device in mesh_devices: |
| if mesh_device is not None: |
| ttnn.synchronize_device(mesh_device) |
| return |
| mesh_device = _target_mesh_device(execution_target) |
| if mesh_device is not None: |
| ttnn.synchronize_device(mesh_device) |
|
|
|
|
| def _concat_host_output(output, cluster_shape): |
| output_tensors = [ttnn.to_torch(tensor) for tensor in ttnn.get_device_tensors(output)] |
| _, columns = cluster_shape |
| mesh_rows = [output_tensors[i : i + columns] for i in range(0, len(output_tensors), columns)] |
| return torch.cat([torch.cat(row, dim=-1) for row in mesh_rows], dim=1) |
|
|
|
|
| def _process_legacy_sampled_tokens(output, batch_size, cluster_shape): |
| torch_output = _concat_host_output(output, cluster_shape) |
| if torch_output.ndim >= 4: |
| if torch_output.shape[2] >= batch_size: |
| return torch_output[0, 0, :batch_size, 0] |
| if torch_output.shape[3] >= batch_size: |
| return torch_output[0, 0, 0, :batch_size] |
| return torch_output.reshape(-1)[:batch_size] |
|
|
|
|
| def _to_host(value, *, blocking): |
| if value is None: |
| return None |
| if isinstance(value, torch.Tensor): |
| return value.cpu() |
| try: |
| return value.cpu(blocking=blocking) |
| except TypeError: |
| return value.cpu() |
|
|
|
|
| def _submit_decode_read(execution_target, decode_output): |
| read_decode_output = getattr(execution_target, "read_decode_output", None) |
| if callable(read_decode_output): |
| return read_decode_output(decode_output, async_read=True) |
|
|
| output, log_probs = _split_output(decode_output) |
| host_output = (_to_host(output, blocking=False), _to_host(log_probs, blocking=False)) |
| return host_output, [ttnn.record_event(_target_mesh_device(execution_target), 0)] |
|
|
|
|
| def _synchronize_read_events(events): |
| if events is None: |
| return |
| if not isinstance(events, (list, tuple, set)): |
| events = [events] |
| for event in events: |
| ttnn.event_synchronize(event) |
|
|
|
|
| def _consume_sampled_output( |
| execution_target, |
| host_output, |
| batch_size, |
| cluster_shape, |
| generated_token_ids, |
| *, |
| process_host_output, |
| ): |
| process_decode_output_host = ( |
| getattr(execution_target, "process_decode_output_host", None) if process_host_output else None |
| ) |
| if callable(process_decode_output_host): |
| tokens, _ = process_decode_output_host(host_output, is_tokens=True) |
| else: |
| tokens, _ = _split_output(host_output) |
| if isinstance(tokens, torch.Tensor): |
| tokens = tokens.view(-1)[:batch_size].detach().cpu() |
| else: |
| tokens = _process_legacy_sampled_tokens(tokens, batch_size, cluster_shape) |
| tokens = tokens.view(-1)[:batch_size].detach().cpu() |
|
|
| for user_id, token in enumerate(tokens.tolist()): |
| generated_token_ids[user_id].append(int(token)) |
|
|
|
|
| def _host_argmax_with_margins(logits: torch.Tensor, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]: |
| """Return host argmax and top1-minus-top2 margins for eval diagnostics.""" |
|
|
| rows = logits[:, -1, :] |
| top2 = torch.topk(rows.float(), k=2, dim=-1) |
| tokens = top2.indices[:, 0].view(-1)[:batch_size].detach().cpu() |
| margins = (top2.values[:, 0] - top2.values[:, 1]).view(-1)[:batch_size].detach().cpu() |
| return tokens, margins |
|
|
|
|
| def run_perf_benchmark( |
| executor, |
| *, |
| tokens: torch.Tensor, |
| kv_cache: list, |
| page_table: torch.Tensor, |
| num_decode_tokens: int = 128, |
| max_batch_size: int = 1, |
| prompt_lens: torch.Tensor | None = None, |
| start_pos: list[int] | None = None, |
| sampling_params=None, |
| prefill_sampling_params=_SAME_SAMPLING_PARAMS, |
| pipeline_readback: bool = False, |
| profiler=None, |
| collect_argmax_diagnostics: bool = False, |
| ) -> PerfBenchmarkResult: |
| """Run the timed prefill and decode loop against a public execution target.""" |
| execution_target = executor |
| mesh_device = _target_mesh_device(execution_target) |
| has_public_readback = callable(getattr(execution_target, "read_decode_output", None)) |
| has_legacy_readback = mesh_device is not None and hasattr(ttnn, "record_event") |
| can_pipeline_readback = ( |
| sampling_params is not None |
| and pipeline_readback |
| and hasattr(ttnn, "event_synchronize") |
| and (has_public_readback or has_legacy_readback) |
| ) |
| if sampling_params is not None and pipeline_readback and not can_pipeline_readback: |
| logger.warning("PIPELINE_READBACK requested, but this execution target does not expose async readback") |
|
|
| batch_size, prompt_len = tokens.shape |
| max_batch_size = max(max_batch_size, batch_size) |
| cluster_shape = _target_cluster_shape(execution_target) |
| prompt_lens = prompt_lens if prompt_lens is not None else torch.tensor([prompt_len] * batch_size) |
| if prefill_sampling_params is _SAME_SAMPLING_PARAMS: |
| prefill_sampling_params = sampling_params |
| prefill_kwargs = dict( |
| |
| |
| |
| page_table=page_table[:batch_size], |
| kv_cache=kv_cache, |
| prompt_lens=prompt_lens, |
| empty_slots=list(range(batch_size)), |
| start_pos=start_pos, |
| ) |
|
|
| compile_tokens = torch.zeros(max_batch_size, prompt_len, dtype=tokens.dtype) |
| compile_tokens[:batch_size] = tokens |
| compile_prompt_lens = torch.zeros(max_batch_size, dtype=prompt_lens.dtype) |
| compile_prompt_lens[:batch_size] = prompt_lens |
| compile_page_table = page_table |
| compile_empty_slots = list(range(batch_size)) |
| if batch_size < max_batch_size: |
| |
| |
| compile_tokens = tokens |
| compile_prompt_lens = prompt_lens |
| compile_page_table = page_table[:batch_size] |
| _compile_prefill_and_decode( |
| execution_target, |
| prefill_tokens=compile_tokens, |
| prefill_page_table=compile_page_table, |
| kv_cache=kv_cache, |
| prompt_lens=compile_prompt_lens, |
| empty_slots=compile_empty_slots, |
| start_pos=start_pos, |
| sampling_params=sampling_params, |
| prefill_sampling_params=prefill_sampling_params, |
| decode_tokens=torch.zeros(max_batch_size, dtype=torch.long), |
| decode_start_pos=torch.full((max_batch_size,), prompt_len, dtype=torch.long), |
| decode_page_table=page_table, |
| ) |
|
|
| _profiler_start(profiler, "inference_prefill") |
| try: |
| start_time = time.perf_counter() |
| prefill_output = execution_target.prefill_forward( |
| tokens, |
| **prefill_kwargs, |
| sampling_params=prefill_sampling_params, |
| ) |
| _synchronize_target(execution_target) |
| prefill_time = time.perf_counter() - start_time |
| finally: |
| _profiler_end(profiler, "inference_prefill") |
|
|
| first_margins = None |
| if collect_argmax_diagnostics and sampling_params is None: |
| first_token, first_margins = _host_argmax_with_margins(prefill_output, batch_size) |
| else: |
| first_token = prefill_output[0] if isinstance(prefill_output, tuple) else torch.argmax(prefill_output, dim=-1) |
| first_token = first_token.view(-1)[:batch_size].detach().cpu() |
| generated_token_ids = [[int(token)] for token in first_token.tolist()] |
| argmax_top2_margins = [[float(margin)] for margin in first_margins.tolist()] if first_margins is not None else None |
|
|
| current_tokens = torch.zeros(max_batch_size, dtype=torch.long) |
| current_tokens[:batch_size] = first_token |
| current_pos = torch.full((max_batch_size,), -1, dtype=torch.long) |
| current_pos[:batch_size] = prompt_lens[:batch_size] |
|
|
| compile_time = None |
| decode_times = [] |
| decode_iteration_times = [] |
| sampled_decode_start = None |
| pending_host_output = None |
| pending_read_events = None |
|
|
| _profiler_start(profiler, "inference_decode") |
| try: |
| for iteration in range(num_decode_tokens): |
| start_time = time.perf_counter() |
| read_from_device = sampling_params is None or not can_pipeline_readback |
| if sampling_params is not None and iteration == 1: |
| sampled_decode_start = start_time |
|
|
| decode_output = execution_target.decode_forward( |
| current_tokens, |
| current_pos, |
| page_table=page_table, |
| kv_cache=kv_cache, |
| read_from_device=read_from_device, |
| sampling_params=sampling_params, |
| reset_batch=iteration == 0, |
| ) |
| output, _ = _split_output(decode_output) |
|
|
| completed_host_output = None |
| if can_pipeline_readback: |
| host_output, read_events = _submit_decode_read(execution_target, decode_output) |
| if pending_read_events is not None: |
| _synchronize_read_events(pending_read_events) |
| completed_host_output = pending_host_output |
| pending_host_output = host_output |
| pending_read_events = read_events |
|
|
| if sampling_params is not None and iteration == 0: |
| _synchronize_target(execution_target) |
| elapsed = time.perf_counter() - start_time |
|
|
| if iteration == 0: |
| compile_time = elapsed |
| else: |
| decode_iteration_times.append(elapsed) |
| if sampling_params is None or can_pipeline_readback: |
| decode_times.append(elapsed) |
|
|
| if completed_host_output is not None: |
| _consume_sampled_output( |
| execution_target, |
| completed_host_output, |
| batch_size, |
| cluster_shape, |
| generated_token_ids, |
| process_host_output=True, |
| ) |
|
|
| if sampling_params is None: |
| if isinstance(output, torch.Tensor) and output.dim() >= 2: |
| if collect_argmax_diagnostics: |
| next_token, next_margins = _host_argmax_with_margins(output, batch_size) |
| else: |
| next_token = torch.argmax(output[:, -1, :], dim=-1) |
| else: |
| if collect_argmax_diagnostics: |
| raise TypeError("argmax margin diagnostics require host decode logits") |
| next_token = output |
| next_token = next_token.view(-1)[:batch_size].detach().cpu() |
| for user_id, token in enumerate(next_token.tolist()): |
| generated_token_ids[user_id].append(int(token)) |
| if argmax_top2_margins is not None: |
| argmax_top2_margins[user_id].append(float(next_margins[user_id])) |
| current_tokens[:batch_size] = next_token |
| elif not can_pipeline_readback: |
| _consume_sampled_output( |
| execution_target, |
| decode_output, |
| batch_size, |
| cluster_shape, |
| generated_token_ids, |
| process_host_output=False, |
| ) |
| current_pos[:batch_size] += 1 |
| finally: |
| _profiler_end(profiler, "inference_decode") |
|
|
| if sampling_params is not None: |
| if pending_read_events is not None: |
| _synchronize_read_events(pending_read_events) |
| _consume_sampled_output( |
| execution_target, |
| pending_host_output, |
| batch_size, |
| cluster_shape, |
| generated_token_ids, |
| process_host_output=True, |
| ) |
| if sampled_decode_start is not None and not can_pipeline_readback: |
| _synchronize_target(execution_target) |
| sampled_decode_time = time.perf_counter() - sampled_decode_start |
| decode_times = [sampled_decode_time / (num_decode_tokens - 1)] * (num_decode_tokens - 1) |
|
|
| return PerfBenchmarkResult( |
| prefill_time_s=prefill_time, |
| compile_decode_time_s=compile_time or 0.0, |
| decode_times_s=decode_times, |
| batch_size=batch_size, |
| num_decode_tokens=num_decode_tokens, |
| generated_token_ids=generated_token_ids, |
| decode_iteration_times_s=decode_iteration_times, |
| argmax_top2_margins=argmax_top2_margins, |
| ) |
|
|
|
|
| def _add_token_ids(target: set[int], value) -> None: |
| if value is None or isinstance(value, bool): |
| return |
| if isinstance(value, int): |
| target.add(int(value)) |
| return |
| if isinstance(value, (list, tuple, set)): |
| for item in value: |
| _add_token_ids(target, item) |
|
|
|
|
| def _stop_token_ids(tokenizer) -> set[int]: |
| stop_ids: set[int] = set() |
| _add_token_ids(stop_ids, getattr(tokenizer, "eos_token_id", None)) |
| _add_token_ids(stop_ids, getattr(tokenizer, "stop_tokens", None)) |
| convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) |
| if callable(convert_tokens_to_ids): |
| eot_id = convert_tokens_to_ids("<|eot_id|>") |
| if isinstance(eot_id, int) and eot_id >= 0: |
| stop_ids.add(eot_id) |
| return stop_ids |
|
|
|
|
| def _truncate_at_stop(token_ids: list[int], stop_ids: set[int]) -> list[int]: |
| for index, token_id in enumerate(token_ids): |
| if token_id in stop_ids: |
| return token_ids[:index] |
| return token_ids |
|
|
|
|
| def assert_no_special_tokens( |
| generated_token_ids: list[list[int]], |
| tokenizer, |
| *, |
| case_name: str = "", |
| is_ci_env: bool | None = None, |
| ) -> None: |
| """Warn locally; fail under CI or ``TT_DEMO_STRICT_SPECIAL_TOKENS=1``.""" |
| if is_ci_env is None: |
| is_ci_env = os.environ.get("CI") == "true" or os.environ.get("TT_DEMO_STRICT_SPECIAL_TOKENS") == "1" |
|
|
| special_ids = set(getattr(tokenizer, "all_special_ids", []) or []) |
| stop_ids = _stop_token_ids(tokenizer) |
| offending_users = 0 |
| for token_ids in generated_token_ids: |
| output_before_stop = _truncate_at_stop(list(token_ids), stop_ids) |
| if any(token_id in special_ids for token_id in output_before_stop): |
| offending_users += 1 |
|
|
| if offending_users == 0: |
| return |
|
|
| prefix = f"[{case_name}] " if case_name else "" |
| message = f"{prefix}model produced special tokens ({offending_users}/{len(generated_token_ids)} users)" |
| logger.warning(message) |
| if is_ci_env: |
| raise AssertionError(message) |
|
|
|
|
| def load_eval_repeat_prompts_batch32() -> list[str]: |
| """The 32 numeric sequence-continuation prompts TTTv1's ci-eval-32 uses (parity).""" |
| path = Path("models/tt_transformers/demo/sample_prompts/eval_repeat_prompts_batch32.json") |
| with open(path) as f: |
| data = json.load(f) |
| return [entry["prompt"] for entry in data] |
|
|
|
|
| def rotate_prompts(all_prompts: list[str], repeat: int) -> list[str]: |
| """Rotate the prompt->slot assignment by ``repeat``: slot j holds prompt (j+repeat)%N.""" |
| n = len(all_prompts) |
| return [all_prompts[(j + repeat) % n] for j in range(n)] |
|
|
|
|
| def eval_page_table_for_repeat(page_table: torch.Tensor, repeat: int, *, mode: str) -> torch.Tensor: |
| """Select physical KV allocation for one rotated eval repeat. |
| |
| ``slot-stable`` is the TTTv1 acceptance geometry: page-table row ``j`` |
| remains at decode slot ``j`` while prompts rotate. ``prompt-stable`` is a |
| diagnostic A/B: rows rotate with prompts, so a prompt retains its original |
| physical KV blocks even as it moves to another decode slot. |
| """ |
|
|
| if mode == "slot-stable": |
| return page_table |
| if mode == "prompt-stable": |
| return torch.roll(page_table, shifts=-repeat, dims=0) |
| raise ValueError(f"unsupported eval page-table mode {mode!r}; use 'slot-stable' or 'prompt-stable'") |
|
|
|
|
| def eval_decode_trace_mode(mode: str) -> str: |
| """Resolve the eval decode execution A/B without changing its default gate.""" |
|
|
| if mode == "traced": |
| return "decode_only" |
| if mode == "eager": |
| return "none" |
| raise ValueError(f"unsupported eval decode mode {mode!r}; use 'traced' or 'eager'") |
|
|
|
|
| def require_canonical_eval_modes_in_ci(environ) -> None: |
| """Prevent diagnostic A/B knobs from replacing a canonical CI gate.""" |
|
|
| if environ.get("CI") != "true": |
| return |
| noncanonical = [] |
| if environ.get("EVAL_DECODE_MODE", "traced") != "traced": |
| noncanonical.append("EVAL_DECODE_MODE") |
| if environ.get("EVAL_PAGE_TABLE_MODE", "slot-stable") != "slot-stable": |
| noncanonical.append("EVAL_PAGE_TABLE_MODE") |
| noncanonical.extend(name for name in ("EVAL_IDENTICAL_PROMPT_INDEX", "EVAL_ACTIVE_BATCH_SIZE") if name in environ) |
| if noncanonical: |
| raise RuntimeError("diagnostic eval modes cannot replace the canonical CI gate: " + ", ".join(noncanonical)) |
|
|
|
|
| def truncate_at_stop(ids: list[int], stop_ids: set[int]) -> list[int]: |
| """Prefix of ``ids`` up to (excluding) the first id in ``stop_ids``.""" |
| out: list[int] = [] |
| for t in ids: |
| if t in stop_ids: |
| break |
| out.append(t) |
| return out |
|
|
|
|
| def hf_stop_ids(tokenizer, hf_model_id: str | None = None) -> set[int]: |
| """Best-effort stop-token id set for an HF ``AutoTokenizer``. |
| |
| Raw HF tokenizers have no ``.stop_tokens`` (that only exists on the TTTv1 wrapped |
| tokenizer). Build the set from ``eos_token_id`` (int|list|None), and — when an |
| ``hf_model_id`` is supplied — also fold in the model's ``generation_config`` eos ids, |
| since chat models (e.g. Llama-3 Instruct) often carry extra eot ids there rather than |
| on ``eos_token_id``. Missing/empty -> empty set (truncation simply runs full length). |
| """ |
| stop: set[int] = set() |
|
|
| def _add(value) -> None: |
| if value is None: |
| return |
| if isinstance(value, bool): |
| return |
| if isinstance(value, int): |
| stop.add(int(value)) |
| elif isinstance(value, (list, tuple, set)): |
| for e in value: |
| _add(e) |
|
|
| _add(getattr(tokenizer, "eos_token_id", None)) |
| |
| |
| _add(getattr(tokenizer, "stop_tokens", None)) |
| if hf_model_id is not None: |
| try: |
| from transformers import GenerationConfig |
|
|
| gen_cfg = GenerationConfig.from_pretrained(hf_model_id) |
| _add(getattr(gen_cfg, "eos_token_id", None)) |
| except Exception as e: |
| logger.debug(f"ci-eval-32: could not read generation_config eos ids for {hf_model_id}: {e}") |
| return stop |
|
|
|
|
| def decode_eval_output(tokenizer, token_ids: list[int], stop_ids: set[int]) -> str: |
| """Return the continuation text used by TTTv1's ci-eval-32 comparison. |
| |
| Different valid BPE segmentations can decode to the same text. The TTTv1 |
| reference stores ``tokenizer.decode(...)`` results and compares those |
| strings, so comparing token-id lists here would make the port stricter than |
| its source workload and report false slot-rotation failures. |
| """ |
| decoded = tokenizer.decode(truncate_at_stop(token_ids, stop_ids)) |
| if not isinstance(decoded, str): |
| raise TypeError(f"tokenizer.decode must return str, got {type(decoded).__name__}") |
| return decoded |
|
|
|
|
| def assert_cross_cardinality_consistency( |
| outputs_by_cardinality: dict[int, dict[str, str]], |
| *, |
| expected_cardinalities: tuple[int, ...] = (1, 2, 4, 32), |
| ) -> None: |
| """Require each fixed request's decoded output to be invariant as batch cardinality grows.""" |
| if tuple(outputs_by_cardinality) != expected_cardinalities: |
| raise AssertionError( |
| f"cross-cardinality experiment expected {expected_cardinalities}, " f"got {tuple(outputs_by_cardinality)}" |
| ) |
| reference: dict[str, tuple[int, str]] = {} |
| for cardinality, outputs in outputs_by_cardinality.items(): |
| if len(outputs) != cardinality: |
| raise AssertionError(f"cardinality {cardinality} returned {len(outputs)} request outputs") |
| for request_id, output in outputs.items(): |
| if request_id in reference: |
| reference_cardinality, reference_output = reference[request_id] |
| if output != reference_output: |
| raise AssertionError( |
| f"request {request_id!r} differs at cardinality {reference_cardinality}->{cardinality}: " |
| f"{reference_output[:120]!r} != {output[:120]!r}" |
| ) |
| else: |
| reference[request_id] = (cardinality, output) |
|
|
|
|
| def assert_cross_batch_consistency( |
| per_repeat_outputs: list[list[str]], |
| *, |
| per_repeat_token_ids: list[list[list[int]]] | None = None, |
| per_repeat_prompt_lens: list[list[int]] | None = None, |
| per_repeat_argmax_margins: list[list[list[float]] | None] | None = None, |
| ) -> None: |
| """Assert decoded prompt-position invariance across repeats. |
| |
| ``per_repeat_outputs[b][u]`` = decoded continuation for slot ``u`` of repeat ``b``. |
| With slot j of repeat b holding prompt (j+b)%N (see ``rotate_prompts``), the same prompt |
| sits at slot (offset+1)%N of repeat b and slot offset of repeat b+1 — so those two |
| outputs must be identical if no per-user state leaks. |
| """ |
| num_batches = len(per_repeat_outputs) |
| assert num_batches >= 2, "cross-batch consistency needs >=2 repeats" |
| n = len(per_repeat_outputs[0]) |
| failed, total = 0, 0 |
| first_failure = None |
| for b in range(num_batches - 1): |
| cur, nxt = per_repeat_outputs[b], per_repeat_outputs[b + 1] |
| for offset in range(n): |
| total += 1 |
| if cur[(offset + 1) % n] != nxt[offset]: |
| failed += 1 |
| if first_failure is None: |
| current_slot = (offset + 1) % n |
| prompt_index = (offset + b + 1) % n |
| token_detail = "" |
| if per_repeat_token_ids is not None: |
| current_tokens = per_repeat_token_ids[b][current_slot] |
| next_tokens = per_repeat_token_ids[b + 1][offset] |
| common_tokens = 0 |
| for current_token, next_token in zip(current_tokens, next_tokens): |
| if current_token != next_token: |
| break |
| common_tokens += 1 |
| current_token = current_tokens[common_tokens] if common_tokens < len(current_tokens) else None |
| next_token = next_tokens[common_tokens] if common_tokens < len(next_tokens) else None |
| token_detail = ( |
| f"; first token divergence at generation step {common_tokens} " |
| f"({current_token!r} != {next_token!r})" |
| ) |
| if per_repeat_argmax_margins is not None: |
| current_repeat_margins = per_repeat_argmax_margins[b] |
| next_repeat_margins = per_repeat_argmax_margins[b + 1] |
| if current_repeat_margins is not None and next_repeat_margins is not None: |
| current_margin = current_repeat_margins[current_slot][common_tokens] |
| next_margin = next_repeat_margins[offset][common_tokens] |
| token_detail += f"; top2 margins {current_margin:.6g} and {next_margin:.6g}" |
| length_detail = "" |
| if per_repeat_prompt_lens is not None: |
| current_length = per_repeat_prompt_lens[b][current_slot] |
| next_length = per_repeat_prompt_lens[b + 1][offset] |
| length_detail = f"; prompt lengths {current_length} and {next_length}" |
| first_failure = ( |
| b, |
| offset, |
| current_slot, |
| prompt_index, |
| cur[current_slot], |
| nxt[offset], |
| token_detail, |
| length_detail, |
| ) |
| assert failed == 0, ( |
| f"ci-eval-32: {failed}/{total} cross-batch consistency checks failed " |
| f"(first at repeat {first_failure[0]} slot {first_failure[2]} -> " |
| f"repeat {first_failure[0] + 1} slot {first_failure[1]}, prompt index {first_failure[3]}" |
| f"{first_failure[6]}{first_failure[7]}; " |
| f"decoded outputs {first_failure[4][:80]!r} != {first_failure[5][:80]!r})" |
| ) |
|
|
|
|
| def assert_within_batch_slot_consistency( |
| decoded_outputs: list[str], |
| *, |
| token_ids: list[list[int]], |
| argmax_margins: list[list[float]] | None, |
| prompt_index: int, |
| ) -> None: |
| """Prove one fixed request is invariant across logical slots in one run.""" |
|
|
| reference = decoded_outputs[0] |
| for slot, decoded in enumerate(decoded_outputs[1:], start=1): |
| if decoded == reference: |
| continue |
| reference_tokens = token_ids[0] |
| slot_tokens = token_ids[slot] |
| common_tokens = 0 |
| for reference_token, slot_token in zip(reference_tokens, slot_tokens): |
| if reference_token != slot_token: |
| break |
| common_tokens += 1 |
| reference_token = reference_tokens[common_tokens] if common_tokens < len(reference_tokens) else None |
| slot_token = slot_tokens[common_tokens] if common_tokens < len(slot_tokens) else None |
| margin_detail = "" |
| if argmax_margins is not None: |
| margin_detail = ( |
| f"; top2 margins {argmax_margins[0][common_tokens]:.6g} " |
| f"and {argmax_margins[slot][common_tokens]:.6g}" |
| ) |
| raise AssertionError( |
| f"ci-eval-32 identical-request diagnostic: prompt index {prompt_index} differs between " |
| f"logical slots 0 and {slot} at generation step {common_tokens} " |
| f"({reference_token!r} != {slot_token!r}){margin_detail}; " |
| f"decoded outputs {reference[:80]!r} != {decoded[:80]!r}" |
| ) |
|
|
|
|
| def run_eval_repeat_batch32( |
| *, |
| make_executor, |
| allocate_kv_cache, |
| page_table: torch.Tensor, |
| prompts: list[str], |
| tokenizer, |
| tokenize_fn, |
| num_decode_tokens: int, |
| max_batch_size: int, |
| sampling_params=None, |
| repeat_batches: int = 3, |
| hf_model_id: str | None = None, |
| first_repeat_profiler=None, |
| page_table_mode: str = "slot-stable", |
| identical_prompt_index: int | None = None, |
| active_batch_size: int | None = None, |
| ) -> PerfBenchmarkResult: |
| """Drive the ci-eval-32 determinism case, building a fresh traced executor per repeat. |
| |
| Each repeat builds its own traced executor (``make_executor()``) and its own zeroed KV |
| cache (``allocate_kv_cache(executor)``), so the rotated batches are fully independent — |
| no shared device or host state can leak across repeats. The executor is cleaned up after |
| each repeat. (The model is bit-deterministic across repeats either way; fresh-per-repeat |
| is simply the cleanest independence guarantee for a determinism test, and the trace |
| recapture cost is negligible at batch-32.) |
| |
| Args: |
| make_executor: Zero-arg callable returning a fresh traced executor |
| (``run_perf_benchmark`` requires traced). Called once per repeat. |
| allocate_kv_cache: Callable(executor) -> fresh zeroed kv_cache bound on that executor. |
| page_table: Fixed contiguous page table (shared across repeats). |
| prompts: The N (=max_batch_size) prompts to rotate (TTTv1 ci-eval-32 numeric prompts; |
| see the module note above re: degenerate-output sensitivity on small models). |
| tokenizer: HF tokenizer (for stop / special ids). |
| tokenize_fn: Callable(list[str]) -> (tokens, prompt_lens). |
| num_decode_tokens: Decode steps per repeat. |
| max_batch_size: Padded batch (== len(prompts) for this fixed-32 case). |
| sampling_params: None -> host argmax (deterministic, mesh-agnostic default). |
| repeat_batches: Number of rotated repeats (TTTv1 uses 3). |
| hf_model_id: Optional, to enrich stop ids from generation_config. |
| first_repeat_profiler: Optional profiler passed only to the first repeat, allowing a caller |
| to emit perf telemetry without changing the three-repeat determinism geometry. |
| page_table_mode: ``slot-stable`` preserves TTTv1 acceptance geometry. |
| ``prompt-stable`` keeps each prompt on the same physical KV blocks |
| as a diagnostic A/B while it rotates through decode slots. |
| identical_prompt_index: Diagnostic-only fixed request replicated into |
| every logical slot. Slot consistency is checked after the first |
| batch, before any repeat-lifecycle comparison. |
| active_batch_size: Diagnostic-only number of leading logical slots to |
| activate while retaining the model's full lane capacity. This is |
| restricted to the identical-request probe so inactive lanes cannot |
| complicate prompt-rotation semantics. |
| """ |
| assert ( |
| len(prompts) == max_batch_size |
| ), f"ci-eval-32 expects len(prompts)==max_batch_size; got {len(prompts)} vs {max_batch_size}" |
| if active_batch_size is not None: |
| if identical_prompt_index is None: |
| raise ValueError("active_batch_size requires identical_prompt_index") |
| if not 1 <= active_batch_size <= max_batch_size: |
| raise ValueError(f"active_batch_size must be in [1, {max_batch_size}]") |
| if identical_prompt_index is not None: |
| if not 0 <= identical_prompt_index < len(prompts): |
| raise ValueError(f"identical_prompt_index must be in [0, {len(prompts) - 1}]") |
| prompts = [prompts[identical_prompt_index]] * (active_batch_size or len(prompts)) |
| stop_ids = hf_stop_ids(tokenizer, hf_model_id) |
| special_ids = set(getattr(tokenizer, "all_special_ids", []) or []) |
| |
| |
| garbage_ids = special_ids - stop_ids |
| logger.info( |
| f"ci-eval-32: repeat_batches={repeat_batches}, N={len(prompts)}, " |
| f"stop_ids={sorted(stop_ids)}, |special_ids|={len(special_ids)}, sampling_params={sampling_params}, " |
| f"page_table_mode={page_table_mode}, identical_prompt_index={identical_prompt_index}, " |
| f"active_batch_size={len(prompts)}" |
| ) |
|
|
| per_repeat: list[list[str]] = [] |
| per_repeat_token_ids: list[list[list[int]]] = [] |
| per_repeat_prompt_lens: list[list[int]] = [] |
| per_repeat_argmax_margins: list[list[list[float]] | None] = [] |
| first_result = None |
| for i in range(repeat_batches): |
| traced_executor = make_executor() |
| try: |
| kv_cache = allocate_kv_cache(traced_executor) |
| rotated = rotate_prompts(prompts, i) |
| tokens, prompt_lens = tokenize_fn(rotated) |
| repeat_page_table = eval_page_table_for_repeat(page_table, i, mode=page_table_mode) |
| result = run_perf_benchmark( |
| traced_executor, |
| tokens=tokens, |
| kv_cache=kv_cache, |
| page_table=repeat_page_table, |
| num_decode_tokens=num_decode_tokens, |
| max_batch_size=max_batch_size, |
| prompt_lens=prompt_lens, |
| sampling_params=sampling_params, |
| profiler=first_repeat_profiler if i == 0 else None, |
| collect_argmax_diagnostics=sampling_params is None, |
| ) |
| finally: |
| traced_executor.cleanup() |
| truncated = [truncate_at_stop(ids, stop_ids) for ids in result.generated_token_ids] |
| for u, ids in enumerate(truncated): |
| bad = set(ids) & garbage_ids |
| assert not bad, f"ci-eval-32: user {u} produced special token(s) {sorted(bad)} mid-stream" |
| decoded = [decode_eval_output(tokenizer, ids, stop_ids) for ids in result.generated_token_ids] |
| per_repeat.append(decoded) |
| per_repeat_token_ids.append(result.generated_token_ids) |
| per_repeat_prompt_lens.append([int(length) for length in prompt_lens]) |
| per_repeat_argmax_margins.append(getattr(result, "argmax_top2_margins", None)) |
| if identical_prompt_index is not None: |
| assert_within_batch_slot_consistency( |
| decoded, |
| token_ids=result.generated_token_ids, |
| argmax_margins=getattr(result, "argmax_top2_margins", None), |
| prompt_index=identical_prompt_index, |
| ) |
| if i == 0: |
| first_result = result |
| logger.info( |
| f"ci-eval-32 repeat {i}: truncated token lengths = {[len(t) for t in truncated]}, " |
| f"decoded character lengths = {[len(text) for text in decoded]}" |
| ) |
|
|
| if identical_prompt_index is None: |
| assert_cross_batch_consistency( |
| per_repeat, |
| per_repeat_token_ids=per_repeat_token_ids, |
| per_repeat_prompt_lens=per_repeat_prompt_lens, |
| per_repeat_argmax_margins=per_repeat_argmax_margins, |
| ) |
| logger.info(f"ci-eval-32: all {(repeat_batches - 1) * len(prompts)} cross-batch consistency checks passed") |
| else: |
| logger.info( |
| f"ci-eval-32 identical-request diagnostic: prompt index {identical_prompt_index} " |
| f"is invariant across all {len(prompts)} logical slots" |
| ) |
| assert first_result is not None |
| return first_result |
|
|
