File size: 14,785 Bytes
a62b6c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Fully resolved construction policy for one prefill runtime."""
from __future__ import annotations
from dataclasses import dataclass, replace
from typing import Any, Callable
from models.common.llm_runtime.config import PageTableLayout
from models.common.llm_runtime.output_reader import OutputReader
from models.common.llm_runtime.prefill.sampling_helpers import _TILE_SIZE
_SUPPORTED_PREFILL_BATCH_SIZES = (1, 2, 4, 8, 16, 32)
@dataclass(frozen=True)
class PrefillRuntimeConfig:
"""Validated collaborators, geometry, and static prefill capabilities."""
model: Any
mesh_device: Any
output_reader: OutputReader
page_table_layout: PageTableLayout # Current geometry; may be replaced once before execution.
page_table_layout_ceiling: PageTableLayout # Construction-time upper bound retained across replacement.
max_batch_size: int
max_prefill_chunk_size: int
supports_batched_prefill: bool | None
disable_batched_prefill: bool
max_prefill_batch_size: int
batched_prefill_batched_extract: bool
cluster_shape: tuple[int, int]
device_sampling_enabled: bool
can_enable_trace: Callable[[int, int], bool]
trace_capture_prime_sequence_lengths: tuple[int, ...]
allow_force_argmax: bool
max_device_top_k: int
sampling_batch_size: int
static_q128_topk_supported: bool
sampling_state_controller: Any
sampling_state: Any
@classmethod
def resolve(
cls,
*,
model: Any,
output_reader: OutputReader,
page_table_layout: PageTableLayout,
max_batch_size: int,
max_prefill_chunk_size: int,
supports_batched_prefill: bool | None = None,
disable_batched_prefill: bool = False,
max_prefill_batch_size: int = 8,
batched_prefill_batched_extract: bool = True,
device_sampling_enabled: bool,
can_enable_trace: Callable[[int, int], bool],
trace_capture_prime_sequence_lengths: tuple[int, ...] = (),
sampling_state_controller: Any = None,
sampling_state: Any = None,
) -> "PrefillRuntimeConfig":
"""Validate construction inputs and derive every static capability."""
_require_positive_int("max_batch_size", max_batch_size)
_require_positive_int("max_prefill_chunk_size", max_prefill_chunk_size)
_require_optional_bool("supports_batched_prefill", supports_batched_prefill)
_require_bool("disable_batched_prefill", disable_batched_prefill)
_require_positive_int("max_prefill_batch_size", max_prefill_batch_size)
_require_supported_prefill_batch_size(max_prefill_batch_size)
_require_bool("batched_prefill_batched_extract", batched_prefill_batched_extract)
if not isinstance(output_reader, OutputReader):
raise TypeError("output_reader must be an OutputReader")
mesh_device = output_reader.mesh_device
try:
resolved_cluster_shape = tuple(int(value) for value in mesh_device.shape)
except (AttributeError, TypeError, ValueError) as error:
raise TypeError("output_reader.mesh_device must provide a two-dimensional shape") from error
if len(resolved_cluster_shape) != 2:
raise ValueError("output_reader.mesh_device shape must contain rows and columns")
for index, value in enumerate(resolved_cluster_shape):
_require_positive_int(f"cluster_shape[{index}]", value)
if not isinstance(device_sampling_enabled, bool):
raise TypeError("device_sampling_enabled must be bool")
if not callable(can_enable_trace):
raise TypeError("can_enable_trace must be callable")
_validate_trace_capture_prime_sequence_lengths(trace_capture_prime_sequence_lengths)
if any(length > max_prefill_chunk_size for length in trace_capture_prime_sequence_lengths):
raise ValueError("trace capture prime sequence lengths cannot exceed max_prefill_chunk_size")
if not isinstance(page_table_layout, PageTableLayout):
raise TypeError("page_table_layout must be a PageTableLayout")
model_mesh = getattr(getattr(model, "config", None), "mesh_device", None)
if model_mesh is None:
model_mesh = getattr(model, "mesh_device", None)
if model_mesh is not mesh_device:
raise ValueError("model and prefill runtime must use the same mesh device")
allow_force_argmax = False
max_device_top_k = 0
sampling_batch_size = int(max_batch_size)
static_q128_topk_supported = False
if device_sampling_enabled:
sampler = getattr(model, "sampling", None)
sampler_config = getattr(sampler, "config", None)
if sampler_config is None:
raise TypeError("device sampling requires model.sampling.config")
allow_force_argmax = getattr(sampler_config, "allow_force_argmax", None)
if not isinstance(allow_force_argmax, bool):
raise TypeError("model.sampling.config.allow_force_argmax must be bool")
max_device_top_k = getattr(sampler_config, "max_top_k", None)
_require_positive_int("model.sampling.config.max_top_k", max_device_top_k)
sampling_batch_size = getattr(sampler_config, "max_batch_size", None)
_require_positive_int("model.sampling.config.max_batch_size", sampling_batch_size)
static_q128_topk_supported = sampling_batch_size >= _TILE_SIZE
return cls(
model=model,
mesh_device=mesh_device,
output_reader=output_reader,
page_table_layout=page_table_layout,
max_batch_size=max_batch_size,
max_prefill_chunk_size=max_prefill_chunk_size,
supports_batched_prefill=supports_batched_prefill,
disable_batched_prefill=disable_batched_prefill,
max_prefill_batch_size=max_prefill_batch_size,
batched_prefill_batched_extract=batched_prefill_batched_extract,
cluster_shape=resolved_cluster_shape,
device_sampling_enabled=device_sampling_enabled,
can_enable_trace=can_enable_trace,
trace_capture_prime_sequence_lengths=trace_capture_prime_sequence_lengths,
allow_force_argmax=allow_force_argmax,
max_device_top_k=max_device_top_k,
sampling_batch_size=sampling_batch_size,
static_q128_topk_supported=static_q128_topk_supported,
sampling_state_controller=sampling_state_controller,
sampling_state=sampling_state,
page_table_layout_ceiling=page_table_layout,
)
def __post_init__(self) -> None:
_require_positive_int("max_batch_size", self.max_batch_size)
_require_positive_int("max_prefill_chunk_size", self.max_prefill_chunk_size)
_require_optional_bool("supports_batched_prefill", self.supports_batched_prefill)
_require_bool("disable_batched_prefill", self.disable_batched_prefill)
_require_positive_int("max_prefill_batch_size", self.max_prefill_batch_size)
_require_supported_prefill_batch_size(self.max_prefill_batch_size)
_require_bool("batched_prefill_batched_extract", self.batched_prefill_batched_extract)
_require_positive_int("sampling_batch_size", self.sampling_batch_size)
if not isinstance(self.output_reader, OutputReader):
raise TypeError("output_reader must be an OutputReader")
if not isinstance(self.page_table_layout, PageTableLayout):
raise TypeError("page_table_layout must be a PageTableLayout")
if not isinstance(self.cluster_shape, tuple):
raise TypeError("cluster_shape must be a tuple")
if len(self.cluster_shape) != 2:
raise ValueError("cluster_shape must contain rows and columns")
for index, value in enumerate(self.cluster_shape):
_require_positive_int(f"cluster_shape[{index}]", value)
if not isinstance(self.device_sampling_enabled, bool):
raise TypeError("device_sampling_enabled must be bool")
if not isinstance(self.allow_force_argmax, bool):
raise TypeError("allow_force_argmax must be bool")
if self.device_sampling_enabled:
_require_positive_int("max_device_top_k", self.max_device_top_k)
elif self.max_device_top_k != 0:
raise ValueError("max_device_top_k must be zero when device sampling is disabled")
if not isinstance(self.static_q128_topk_supported, bool):
raise TypeError("static_q128_topk_supported must be bool")
_validate_trace_capture_prime_sequence_lengths(self.trace_capture_prime_sequence_lengths)
if not isinstance(self.page_table_layout_ceiling, PageTableLayout):
raise TypeError("page_table_layout_ceiling must be a PageTableLayout")
if self.page_table_layout.block_size != self.page_table_layout_ceiling.block_size:
raise ValueError("page_table_layout_ceiling cannot change block_size")
if self.page_table_layout.raw_capacity_width > self.page_table_layout_ceiling.raw_capacity_width:
raise ValueError("page_table_layout_ceiling must cover page_table_layout capacity")
if (
self.page_table_layout.prefill_width > self.page_table_layout_ceiling.prefill_width
or self.page_table_layout.decode_width > self.page_table_layout_ceiling.decode_width
):
raise ValueError("page_table_layout_ceiling must cover canonical page-table geometry")
if not callable(self.can_enable_trace):
raise TypeError("can_enable_trace must be callable")
if self.mesh_device is None:
raise ValueError("mesh_device is required")
model_mesh = getattr(getattr(self.model, "config", None), "mesh_device", None)
if model_mesh is None:
model_mesh = getattr(self.model, "mesh_device", None)
if model_mesh is not self.mesh_device:
raise ValueError("model and prefill runtime must use the same mesh device")
if self.output_reader.mesh_device is not self.mesh_device:
raise ValueError("output_reader and prefill runtime must use the same mesh device")
expected_argmax = False
expected_max_top_k = 0
expected_sampling_batch_size = self.max_batch_size
expected_q128 = False
if self.device_sampling_enabled:
sampler_config = getattr(getattr(self.model, "sampling", None), "config", None)
if sampler_config is None:
raise TypeError("device sampling requires model.sampling.config")
expected_argmax = getattr(sampler_config, "allow_force_argmax", None)
if not isinstance(expected_argmax, bool):
raise TypeError("model.sampling.config.allow_force_argmax must be bool")
expected_max_top_k = getattr(sampler_config, "max_top_k", None)
_require_positive_int("model.sampling.config.max_top_k", expected_max_top_k)
expected_sampling_batch_size = getattr(sampler_config, "max_batch_size", None)
_require_positive_int("model.sampling.config.max_batch_size", expected_sampling_batch_size)
expected_q128 = expected_sampling_batch_size >= _TILE_SIZE
if self.allow_force_argmax is not expected_argmax:
raise ValueError("allow_force_argmax must match the resolved sampler capability")
if self.max_device_top_k != expected_max_top_k:
raise ValueError("max_device_top_k must match the resolved sampler capability")
if self.sampling_batch_size != expected_sampling_batch_size:
raise ValueError("sampling_batch_size must match the resolved sampler capacity")
if self.static_q128_topk_supported is not expected_q128:
raise ValueError("static_q128_topk_supported must match the resolved sampler capacity")
if (self.sampling_state_controller is None) != (self.sampling_state is None):
raise ValueError("sampling_state_controller and sampling_state must be supplied together")
if self.sampling_state_controller is not None:
if getattr(self.sampling_state_controller, "sampling", None) is not getattr(self.model, "sampling", None):
raise ValueError("sampling state controller must borrow model.sampling")
if not callable(getattr(self.sampling_state_controller, "admit", None)):
raise TypeError("sampling state controller must provide admit()")
def with_page_table_layout(self, layout: PageTableLayout) -> "PrefillRuntimeConfig":
"""Return the same resolved policy with a smaller final KV geometry."""
if not isinstance(layout, PageTableLayout):
raise TypeError("layout must be a PageTableLayout")
if layout.block_size != self.page_table_layout.block_size:
raise ValueError("page-table layout replacement cannot change block_size")
if layout.raw_capacity_width > self.page_table_layout_ceiling.raw_capacity_width:
raise ValueError("page-table layout replacement cannot exceed the construction-time capacity ceiling")
if (
layout.prefill_width > self.page_table_layout_ceiling.prefill_width
or layout.decode_width > self.page_table_layout_ceiling.decode_width
):
raise ValueError("page-table layout replacement cannot expand canonical geometry")
return replace(self, page_table_layout=layout)
def _require_positive_int(name: str, value: Any) -> None:
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise ValueError(f"{name} must be a positive integer")
def _require_bool(name: str, value: Any) -> None:
if not isinstance(value, bool):
raise TypeError(f"{name} must be bool")
def _require_optional_bool(name: str, value: Any) -> None:
if value is not None:
_require_bool(name, value)
def _require_supported_prefill_batch_size(value: int) -> None:
if value not in _SUPPORTED_PREFILL_BATCH_SIZES:
raise ValueError(f"max_prefill_batch_size must be one of {_SUPPORTED_PREFILL_BATCH_SIZES}")
def _validate_trace_capture_prime_sequence_lengths(values: Any) -> None:
if not isinstance(values, tuple):
raise TypeError("trace_capture_prime_sequence_lengths must be a tuple")
if len(set(values)) != len(values):
raise ValueError("trace_capture_prime_sequence_lengths must not contain duplicates")
for index, value in enumerate(values):
_require_positive_int(f"trace_capture_prime_sequence_lengths[{index}]", value)
|