| |
| |
|
|
| """Static configuration and resolved geometry for the LLM runtime toolkit.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Literal |
|
|
| import ttnn |
|
|
| TraceMode = Literal["none", "decode_only", "all"] |
|
|
| _PAGE_TABLE_WIDTH_ALIGNMENT = 8 |
|
|
|
|
| @dataclass(frozen=True) |
| class TraceConfig: |
| """Select which already-compiled graphs receive trace artifacts.""" |
|
|
| mode: TraceMode = "decode_only" |
|
|
| def __post_init__(self) -> None: |
| if self.mode not in ("none", "decode_only", "all"): |
| raise ValueError(f"Unsupported trace mode: {self.mode!r}") |
|
|
| @property |
| def prefill_enabled(self) -> bool: |
| return self.mode == "all" |
|
|
| @property |
| def decode_enabled(self) -> bool: |
| return self.mode in ("decode_only", "all") |
|
|
|
|
| @dataclass(frozen=True) |
| class WarmupConfig: |
| """Select graph coverage while retaining model-derived defaults. |
| |
| ``prefill_seq_lens=None`` asks the executor to use the model runtime's |
| supported trace lengths, falling back to 128. Logits coverage is always |
| required. Forced-top-k prefill and greedy decode coverage are added when |
| dynamic warmup enables device sampling. Additional top-k decode coverage is |
| explicitly opt in. |
| """ |
|
|
| prefill_seq_lens: tuple[int, ...] | None = None |
| prefill_batch_sizes: tuple[int, ...] = (1, 2, 4, 8, 16, 32) |
| include_decode_top_k: bool = False |
|
|
| def __post_init__(self) -> None: |
| if self.prefill_seq_lens is not None: |
| self._validate_positive_tuple("prefill_seq_lens", self.prefill_seq_lens) |
| self._validate_positive_tuple("prefill_batch_sizes", self.prefill_batch_sizes) |
|
|
| @staticmethod |
| def _validate_positive_tuple(name: str, values: tuple[int, ...]) -> None: |
| if not isinstance(values, tuple): |
| raise TypeError(f"{name} must be a tuple so WarmupConfig remains immutable") |
| if not values: |
| raise ValueError(f"{name} must not be empty") |
| if any(not isinstance(value, int) or isinstance(value, bool) or value <= 0 for value in values): |
| raise ValueError(f"{name} values must be positive integers") |
| if len(set(values)) != len(values): |
| raise ValueError(f"{name} values must be unique") |
|
|
|
|
| @dataclass(frozen=True) |
| class PagedKVCacheConfig: |
| """Static paged-KV policy plus an optional resolved physical capacity.""" |
|
|
| block_size: int |
| max_num_blocks: int |
| dtype: ttnn.DataType |
| memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG |
| num_blocks: int | None = None |
|
|
| def __post_init__(self) -> None: |
| if not isinstance(self.block_size, int) or isinstance(self.block_size, bool) or self.block_size <= 0: |
| raise ValueError("block_size must be a positive integer") |
| if ( |
| not isinstance(self.max_num_blocks, int) |
| or isinstance(self.max_num_blocks, bool) |
| or self.max_num_blocks <= 0 |
| ): |
| raise ValueError("max_num_blocks must be a positive integer") |
| if self.num_blocks is not None: |
| if not isinstance(self.num_blocks, int) or isinstance(self.num_blocks, bool) or self.num_blocks <= 0: |
| raise ValueError("num_blocks must be a positive integer when resolved") |
| if self.num_blocks > self.max_num_blocks: |
| raise ValueError(f"num_blocks ({self.num_blocks}) exceeds max_num_blocks ({self.max_num_blocks})") |
|
|
| def is_resolved(self) -> bool: |
| return self.num_blocks is not None |
|
|
|
|
| @dataclass(frozen=True) |
| class PageTableLayout: |
| """Resolved page-table geometry shared by prefill, decode, and warmup.""" |
|
|
| block_size: int |
| raw_capacity_width: int |
| prefill_width: int |
| decode_width: int |
|
|
| @classmethod |
| def resolve( |
| cls, |
| *, |
| block_size: int, |
| model_max_sequence_length: int, |
| physical_num_blocks: int, |
| max_prefill_chunk_size: int, |
| ) -> "PageTableLayout": |
| values = { |
| "block_size": block_size, |
| "model_max_sequence_length": model_max_sequence_length, |
| "physical_num_blocks": physical_num_blocks, |
| "max_prefill_chunk_size": max_prefill_chunk_size, |
| } |
| for name, value in values.items(): |
| if not isinstance(value, int) or isinstance(value, bool) or value <= 0: |
| raise ValueError(f"{name} must be a positive integer") |
|
|
| model_width = _ceil_div(model_max_sequence_length, block_size) |
| raw_width = min(model_width, physical_num_blocks) |
| decode_width = _round_up(raw_width, _PAGE_TABLE_WIDTH_ALIGNMENT) |
| padding_blocks = _ceil_div(max_prefill_chunk_size - 1, block_size) |
| prefill_width = _round_up(raw_width + padding_blocks, _PAGE_TABLE_WIDTH_ALIGNMENT) |
| return cls( |
| block_size=block_size, |
| raw_capacity_width=raw_width, |
| prefill_width=prefill_width, |
| decode_width=decode_width, |
| ) |
|
|
| def __post_init__(self) -> None: |
| if self.block_size <= 0 or self.raw_capacity_width <= 0: |
| raise ValueError("page-table block size and capacity width must be positive") |
| if self.prefill_width < self.raw_capacity_width or self.decode_width < self.raw_capacity_width: |
| raise ValueError("canonical page-table widths cannot be smaller than raw capacity") |
| if self.prefill_width % _PAGE_TABLE_WIDTH_ALIGNMENT or self.decode_width % _PAGE_TABLE_WIDTH_ALIGNMENT: |
| raise ValueError("canonical page-table widths must satisfy alignment") |
|
|
|
|
| def _ceil_div(value: int, divisor: int) -> int: |
| return (value + divisor - 1) // divisor |
|
|
|
|
| def _round_up(value: int, alignment: int) -> int: |
| return _ceil_div(value, alignment) * alignment |
|
|
