# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import dataclasses import inspect from types import SimpleNamespace from typing import Any import pytest import torch import models.common.llm_runtime.decode as decode_module import ttnn from models.common.llm_runtime.config import PageTableLayout from models.common.llm_runtime.decode import ( DecodeDeviceInputs, DecodePersistentInputs, DecodeProgramSignature, DecodeRuntime, DecodeRuntimeConfig, DecodeTraceSignature, InvocationResult, ) from models.common.llm_runtime.output_reader import OutputReader, PendingRead from models.common.modules.sampling.params import PreparedSamplingParams from models.common.sampling.sampling_params import SamplingParams class FakeMesh: shape = (1, 1) class FakeSampling: def __init__(self, seed_buffer=None): self.config = SimpleNamespace(allow_force_argmax=True, max_batch_size=2, max_top_k=32, seeds=seed_buffer) def decode_forward( self, logits, *, k=None, p=None, temp=None, seeds=None, tt_out_tok=None, enable_log_probs=False, ): return logits, None class FakeRope: def get_rot_idxs(self, positions, *, on_host): assert on_host return ("rotary", positions.clone()) def get_rot_mats(self, rotary_indices): return ("cos", "sin") class FakeModel: def __init__(self, seed_buffer=None): self.config = SimpleNamespace(max_batch_size=2) self.sampling = FakeSampling(seed_buffer) self.rope_setup = FakeRope() self.vocab_size = 8 self.num_devices = 1 def iter_executor_named_modules(self): return iter(()) def increment_positions(self, positions, rotary_indices): return None class FakeLazySeedBuffer: def __init__(self): self.source = torch.arange(2, dtype=torch.int64) self._value = object() self.updates = [] def update(self, source): self.source = source self.updates.append(source.clone()) def get_device_buffer(self): return self._value def make_runtime(*, sampling=True, force_greedy_top_k=False, seed_buffer=None): mesh = FakeMesh() model = FakeModel(seed_buffer) config = DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=2, page_table_layout=page_table_layout(), device_sampling_enabled=sampling, force_greedy_top_k=force_greedy_top_k, ) return DecodeRuntime(config) def page_table_layout(*, raw_width=8, block_size=32): return PageTableLayout( block_size=block_size, raw_capacity_width=raw_width, prefill_width=((raw_width + 7) // 8) * 8, decode_width=((raw_width + 7) // 8) * 8, ) def greedy_sampling(): return SamplingParams(temperature=[0.0, 0.0], top_k=[1, 1], top_p=[1.0, 1.0]) def test_prepare_uses_resolved_sampler_capacity_and_neutral_inactive_rows(): prepared = prepare( make_runtime(), positions=(0, -1), sampling_params=SamplingParams(temperature=1.0, top_k=32, top_p=0.08), ).prepared_sampling assert isinstance(prepared, PreparedSamplingParams) assert prepared.batch_size == 2 assert prepared.active_rows == 1 assert prepared.active_mask == (True, False) assert prepared.top_k == (32, 1) assert prepared.top_p == pytest.approx((0.08, 0.0)) assert prepared.temperature == (1.0, 1.0) assert prepared.row_paths == ("topk", "inactive") def test_prepare_accepts_vector_tensor_fields_for_full_lane(): prepared = prepare( make_runtime(), positions=(0, 0), sampling_params=SamplingParams( temperature=torch.ones(2), top_k=torch.full((2,), 32, dtype=torch.int32), top_p=torch.full((2,), 0.08), ), ).prepared_sampling assert isinstance(prepared, PreparedSamplingParams) assert prepared.active_mask == (True, True) assert prepared.top_k == (32, 32) assert prepared.top_p == pytest.approx((0.08, 0.08)) assert prepared.temperature == (1.0, 1.0) def prepare( runtime, *, positions=(0, -1), page_table=None, sampling_params=None, prompt_tokens=None, output_tokens=None, slot_remap=None, reset=False, ): if page_table is None: page_table = torch.tensor([[3, 4, 5], [6, 7, 8]], dtype=torch.int32) return runtime.prepare( torch.tensor([11, 0]), torch.tensor(positions), page_table, sampling_params=sampling_params, prompt_tokens=prompt_tokens, output_tokens=output_tokens, slot_remap=slot_remap, reset_batch=reset, ) def test_prepare_places_only_start_pos_active_rows_and_neutralizes_gap_sentinels(): runtime = make_runtime() sampling = SamplingParams( temperature=[0.8, 0.8], top_k=[999, 7], top_p=[0.9, 0.8], seed=[-1, -1], enable_log_probs=[False, False], num_logprobs=[-2, -2], ) prepared = prepare( runtime, positions=(-1, 4), sampling_params=sampling, prompt_tokens=torch.tensor([[10, -1], [20, 21]]), output_tokens=[[30, -1], [40, 41]], slot_remap=torch.tensor([1, 0]), reset=True, ).prepared_sampling assert prepared is not None assert prepared.active_mask == (False, True) assert prepared.row_paths == ("inactive", "topk") assert prepared.top_k == (1, 7) assert prepared.top_p == pytest.approx((0.0, 0.8)) assert prepared.seeds == (None, None) assert prepared.enable_log_probs == (False, False) assert prepared.num_logprobs == (0, 0) assert prepared.prompt_tokens.tolist() == [[-1, -1], [20, 21]] assert prepared.output_tokens == [[-1, -1], [40, 41]] assert prepared.slot_remap.tolist() == [1, 0] def seeded_sampling(seed0, seed1=None): return SamplingParams( temperature=[0.8, 0.8], top_k=[32, 32], top_p=[0.95, 0.95], seed=[seed0, seed1], ) def stochastic_sampling(seed=None): return SamplingParams( temperature=[0.8, 0.8], top_k=[32, 32], top_p=[0.95, 0.95], seed=seed, ) def make_four_slot_seed_runtime(): seed_buffer = FakeLazySeedBuffer() seed_buffer.source = torch.arange(4, dtype=torch.int64) model = FakeModel(seed_buffer) model.config.max_batch_size = 4 model.sampling.config.max_batch_size = 4 runtime = DecodeRuntime( DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(FakeMesh()), lane_capacity=4, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) ) return runtime, seed_buffer def four_slot_sampled_warmup(runtime): return runtime.prepare( torch.zeros(4, dtype=torch.long), torch.zeros(4, dtype=torch.long), torch.zeros((4, 8), dtype=torch.int32), sampling_params=SamplingParams( temperature=torch.ones(4), top_k=torch.full((4,), 32, dtype=torch.int32), top_p=torch.full((4,), 0.08), seed=[11, 22, 33, 44], ), ) def _stub_compile_only_decode(runtime, monkeypatch, run_body): monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: "host") monkeypatch.setattr( runtime, "_stage_inputs_and_kpt", lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None), ) monkeypatch.setattr(runtime, "_run_body", run_body) def test_compile_only_sampled_decode_temporarily_admits_and_resets_fallback_seed_slots(monkeypatch, expect_error): runtime, seed_buffer = make_four_slot_seed_runtime() defaults = seed_buffer.source.clone() prepared = four_slot_sampled_warmup(runtime) during_compile = [] def run_body(*args, **kwargs): during_compile.append(runtime._seed_state.snapshot()) assert not kwargs["count_tokens"] assert not kwargs["advance_seeds"] return object() _stub_compile_only_decode(runtime, monkeypatch, run_body) runtime.invoke(prepared, count_tokens=False) assert during_compile[0].active_slots == (0, 1, 2, 3) assert not during_compile[0].buffer_is_default reset = runtime._seed_state.snapshot() assert reset.active_slots == () assert reset.buffer_is_default assert torch.equal(seed_buffer.source, defaults) with expect_error(RuntimeError, "reset_batch=True"): runtime._refresh_sampling_seeds(prepared) def test_compile_only_sampled_decode_resets_fallback_seed_slots_after_failure(monkeypatch, expect_error): runtime, seed_buffer = make_four_slot_seed_runtime() defaults = seed_buffer.source.clone() prepared = four_slot_sampled_warmup(runtime) def run_body(*args, **kwargs): assert runtime._seed_state.snapshot().active_slots == (0, 1, 2, 3) raise RuntimeError("compile boom") _stub_compile_only_decode(runtime, monkeypatch, run_body) monkeypatch.setattr(runtime, "_release_or_retain_transient", lambda owned: []) with expect_error(RuntimeError, "compile boom"): runtime.invoke(prepared, count_tokens=False) reset = runtime._seed_state.snapshot() assert reset.active_slots == () assert reset.buffer_is_default assert torch.equal(seed_buffer.source, defaults) def test_runtime_seed_same_request_and_absolute_position_are_cardinality_independent(): first_buffer = FakeLazySeedBuffer() first = make_runtime(seed_buffer=first_buffer) first_prepared = prepare(first, positions=(41, -1), sampling_params=seeded_sampling(1234), reset=True) first._refresh_sampling_seeds(first_prepared) remapped_buffer = FakeLazySeedBuffer() remapped = make_runtime(seed_buffer=remapped_buffer) remapped_prepared = prepare( remapped, positions=(-1, 41), sampling_params=seeded_sampling(None, 1234), reset=True, ) remapped._refresh_sampling_seeds(remapped_prepared) assert int(first_buffer.updates[-1][0]) == int(remapped_buffer.updates[-1][1]) assert first._seed_state.snapshot().active == (True, False) assert remapped._seed_state.snapshot().active == (False, True) @pytest.mark.parametrize("seed", [1234, torch.tensor(1234)]) def test_runtime_scalar_seed_belongs_to_one_request_and_is_not_broadcast(seed): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) sampling_params = SamplingParams( temperature=[0.8, 0.8], top_k=[32, 32], top_p=[0.95, 0.95], seed=seed, ) prepared = prepare(runtime, positions=(17, 17), sampling_params=sampling_params, reset=True) runtime._refresh_sampling_seeds(prepared) assert prepared.prepared_sampling.seeds == (1234, None) snapshot = runtime._seed_state.snapshot() assert snapshot.request_seeds == (1234, None) assert snapshot.active == (True, True) assert snapshot.current_device_seeds[0] is not None assert snapshot.current_device_seeds[1] is not None @pytest.mark.parametrize("seed", [[111, 222], torch.tensor([111, 222])]) def test_runtime_vector_seed_remains_slot_indexed(seed): runtime = make_runtime(seed_buffer=FakeLazySeedBuffer()) sampling_params = SamplingParams( temperature=[0.8, 0.8], top_k=[32, 32], top_p=[0.95, 0.95], seed=seed, ) prepared = prepare(runtime, positions=(5, 5), sampling_params=sampling_params, reset=True) assert prepared.prepared_sampling.seeds == (111, 222) def test_runtime_simultaneous_equal_request_seeds_share_vllm_stream(): runtime = make_runtime(seed_buffer=FakeLazySeedBuffer()) prepared = prepare( runtime, positions=(5, 5), sampling_params=seeded_sampling(77, 77), reset=True, ) runtime._refresh_sampling_seeds(prepared) snapshot = runtime._seed_state.snapshot() assert runtime._seed_manager.salt_duplicate_seeds is False assert snapshot.request_seeds == (77, 77) assert snapshot.salts == (0, 0) assert snapshot.current_device_seeds[0] == snapshot.current_device_seeds[1] def test_runtime_vllm_uniform_seed_is_deterministic_across_concurrent_slots(): seed_buffer = FakeLazySeedBuffer() seed_buffer.source = torch.arange(32, dtype=torch.int64) model = FakeModel(seed_buffer) model.config.max_batch_size = 32 model.sampling.config.max_batch_size = 32 runtime = DecodeRuntime( DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(FakeMesh()), lane_capacity=32, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) ) prepared = runtime.prepare( torch.zeros(32, dtype=torch.long), torch.zeros(32, dtype=torch.long), torch.zeros((32, 8), dtype=torch.int32), sampling_params=SamplingParams( temperature=torch.ones(32), top_k=torch.full((32,), 32, dtype=torch.int32), top_p=torch.full((32,), 0.95), seed=[1234] * 32, ), reset_batch=True, ) runtime._refresh_sampling_seeds(prepared) snapshot = runtime._seed_state.snapshot() assert runtime._seed_manager.salt_duplicate_seeds is False assert snapshot.salts == (0,) * 32 assert len(set(snapshot.current_device_seeds)) == 1 def test_runtime_slot_remap_moves_complete_seed_stream_before_refresh(): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) initial = prepare( runtime, positions=(5, -1), sampling_params=seeded_sampling(77), reset=True, ) runtime._refresh_sampling_seeds(initial) original_device_seed = int(seed_buffer.updates[-1][0]) original_state = runtime._seed_state.snapshot() moved = prepare( runtime, positions=(-1, 5), sampling_params=seeded_sampling(None, 77), slot_remap=torch.tensor([0, 0]), reset=False, ) runtime._refresh_sampling_seeds(moved) state = runtime._seed_state.snapshot() assert state.active == (False, True) assert state.request_seeds == (None, 77) assert state.token_counters[1] == original_state.token_counters[0] assert int(seed_buffer.updates[-1][1]) == original_device_seed def test_runtime_seed_changes_with_request_seed_and_decode_position(): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) runtime._refresh_sampling_seeds( prepare(runtime, positions=(7, -1), sampling_params=seeded_sampling(101), reset=True) ) position_7 = int(seed_buffer.updates[-1][0]) runtime._refresh_sampling_seeds( prepare(runtime, positions=(8, -1), sampling_params=seeded_sampling(101), reset=False) ) position_8 = int(seed_buffer.updates[-1][0]) runtime._refresh_sampling_seeds( prepare(runtime, positions=(7, -1), sampling_params=seeded_sampling(202), reset=True) ) different_request = int(seed_buffer.updates[-1][0]) assert len({position_7, position_8, different_request}) == 3 def test_runtime_explicit_seed_absolute_position_is_stable_across_reset_boundaries(): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) request = seeded_sampling(909) runtime._refresh_sampling_seeds(prepare(runtime, positions=(19, -1), sampling_params=request, reset=True)) original = seed_buffer.updates[-1].clone() runtime._refresh_sampling_seeds(prepare(runtime, positions=(20, -1), sampling_params=request, reset=False)) continued = seed_buffer.updates[-1].clone() runtime._refresh_sampling_seeds(prepare(runtime, positions=(19, -1), sampling_params=request, reset=True)) restarted = seed_buffer.updates[-1].clone() runtime._refresh_sampling_seeds(prepare(runtime, positions=(20, -1), sampling_params=request, reset=True)) resumed = seed_buffer.updates[-1].clone() assert torch.equal(original, restarted) assert torch.equal(continued, resumed) def test_runtime_seed_refreshes_before_eager_model_invocation(monkeypatch): events = [] seed_buffer = FakeLazySeedBuffer() original_update = seed_buffer.update def record_update(source): events.append("seed") original_update(source) seed_buffer.update = record_update runtime = make_runtime(seed_buffer=seed_buffer) prepared = prepare(runtime, positions=(3, -1), sampling_params=seeded_sampling(77), reset=True) monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: object()) monkeypatch.setattr( runtime, "_stage_inputs_and_kpt", lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None), ) def run_body(*args, **kwargs): events.append("invoke") return object() monkeypatch.setattr(runtime, "_run_body", run_body) runtime.invoke(prepared) assert events[-1] == "invoke" assert events[:-1] assert set(events[:-1]) == {"seed"} def test_runtime_trace_captures_stable_seed_handle_and_refreshes_before_replay(monkeypatch): events = [] seed_buffer = FakeLazySeedBuffer() original_update = seed_buffer.update def record_update(source): events.append("seed") original_update(source) seed_buffer.update = record_update runtime = make_runtime(seed_buffer=seed_buffer) prepared = prepare(runtime, positions=(12, -1), sampling_params=seeded_sampling(88), reset=True) monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda prepared: object()) monkeypatch.setattr( runtime, "_stage_inputs_and_kpt", lambda host, prepared: (DecodeDeviceInputs(None, None, None, None), None), ) persistent = runtime.capture_plan(prepared).prepare_inputs() assert persistent.seed_buffer is seed_buffer.get_device_buffer() sampling = prepared.prepared_sampling assert sampling is not None persistent = dataclasses.replace( persistent, kpt_signature=[(sampling.top_k, sampling.top_p, sampling.temperature)], ) runtime.refresh_trace(persistent, prepared, SimpleNamespace(full=False, page_table=False)) events.append("replay") assert events[-1] == "replay" assert events[:-1] assert set(events[:-1]) == {"seed"} def test_runtime_seed_handling_does_not_mutate_sampling_params(): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) sampling_params = seeded_sampling(11, 22) before = dataclasses.asdict(sampling_params) prepared = prepare(runtime, positions=(4, 9), sampling_params=sampling_params, reset=True) runtime._refresh_sampling_seeds(prepared) assert dataclasses.asdict(sampling_params) == before def test_runtime_unseeded_stream_varies_and_same_absolute_position_is_idempotent(): seed_buffer = FakeLazySeedBuffer() defaults = seed_buffer.source.clone() runtime = make_runtime(seed_buffer=seed_buffer) unseeded = prepare(runtime, positions=(5, -1), sampling_params=stochastic_sampling(), reset=True) runtime._refresh_sampling_seeds(unseeded) first = seed_buffer.updates[-1].clone() first_state = runtime._seed_state.snapshot() assert int(first[0]) != int(defaults[0]) assert int(first[1]) == int(defaults[1]) runtime._refresh_sampling_seeds(dataclasses.replace(unseeded, reset_batch=False)) repeated = seed_buffer.updates[-1].clone() assert torch.equal(repeated, first) advanced = prepare( runtime, positions=(6, -1), sampling_params=stochastic_sampling(), reset=False, ) runtime._refresh_sampling_seeds(advanced) snapshot = runtime._seed_state.snapshot() assert snapshot.request_seeds == (None, None) assert snapshot.active == (True, False) assert snapshot.token_counters[0] == 2 assert snapshot.unseeded_rng_states[0] != first_state.unseeded_rng_states[0] def test_runtime_seed_change_requires_reset_and_preserves_unseeded_survivor(expect_error): seed_buffer = FakeLazySeedBuffer() runtime = make_runtime(seed_buffer=seed_buffer) unseeded = prepare( runtime, positions=(10, 10), sampling_params=stochastic_sampling(seed=[None, None]), reset=True, ) runtime._refresh_sampling_seeds(unseeded) initial = runtime._seed_state.snapshot() mixed = prepare( runtime, positions=(11, 11), sampling_params=stochastic_sampling(seed=[None, 42]), reset=False, ) with expect_error(RuntimeError, "reset_batch=True"): runtime._refresh_sampling_seeds(mixed) runtime._refresh_sampling_seeds(dataclasses.replace(mixed, reset_batch=True)) admitted = runtime._seed_state.snapshot() assert admitted.request_seeds == (None, 42) assert admitted.active == (True, True) assert admitted.token_counters[0] == initial.token_counters[0] + 1 continued = dataclasses.replace(mixed, start_pos=torch.tensor([12, 12])) runtime._refresh_sampling_seeds(continued) continued_state = runtime._seed_state.snapshot() assert continued_state.token_counters[0] == admitted.token_counters[0] + 1 assert continued_state.request_seeds == (None, 42) def test_runtime_inactive_peer_is_cleaned_up_without_readmitting_survivor(): seed_buffer = FakeLazySeedBuffer() defaults = seed_buffer.source.clone() runtime = make_runtime(seed_buffer=seed_buffer) initial = prepare( runtime, positions=(20, 20), sampling_params=stochastic_sampling(seed=[None, 42]), reset=True, ) runtime._refresh_sampling_seeds(initial) initial_state = runtime._seed_state.snapshot() peer_left = prepare( runtime, positions=(21, -1), sampling_params=stochastic_sampling(seed=[None, 42]), reset=False, ) runtime._refresh_sampling_seeds(peer_left) state = runtime._seed_state.snapshot() assert state.active == (True, False) assert state.request_seeds == (None, None) assert state.token_counters[0] == initial_state.token_counters[0] + 1 assert int(seed_buffer.updates[-1][1]) == int(defaults[1]) @pytest.mark.parametrize( ("method", "expected"), ( ( DecodeRuntime.prepare, ( ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("tokens", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("start_pos", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("page_table", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("sampling_params", inspect.Parameter.KEYWORD_ONLY, None), ("prompt_tokens", inspect.Parameter.KEYWORD_ONLY, None), ("output_tokens", inspect.Parameter.KEYWORD_ONLY, None), ("slot_remap", inspect.Parameter.KEYWORD_ONLY, None), ("reset_batch", inspect.Parameter.KEYWORD_ONLY, False), ), ), ( DecodeRuntime.read_decode_output, ( ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("async_read", inspect.Parameter.KEYWORD_ONLY, False), ), ), ( DecodeRuntime.process_decode_output_host, ( ("self", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("tt_out", inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.empty), ("is_tokens", inspect.Parameter.KEYWORD_ONLY, False), ), ), ), ) def test_runtime_api_signatures_are_exact(method, expected): parameters = inspect.signature(method).parameters.values() assert tuple((parameter.name, parameter.kind, parameter.default) for parameter in parameters) == expected def test_runtime_api_preserves_positional_prefixes_and_rejects_extra_fields(expect_error): cases = ( (DecodeRuntime.prepare, ("tokens", "start_pos", "page_table"), "sampling_params"), (DecodeRuntime.read_decode_output, ("tt_out",), "async_read"), (DecodeRuntime.process_decode_output_host, ("tt_out",), "is_tokens"), ) for method, positional_prefix, keyword_name in cases: signature = inspect.signature(method) signature.bind(None, *positional_prefix, **{keyword_name: True}) with expect_error(TypeError, "too many positional arguments"): signature.bind(None, *positional_prefix, True) with expect_error(TypeError, "unexpected keyword argument 'unknown'"): signature.bind(None, *positional_prefix, unknown=True) assert inspect.get_annotations(DecodeRuntime.prepare, eval_str=True)["sampling_params"] is Any def test_config_resolves_canonical_static_capabilities_and_is_frozen(expect_error): mesh = FakeMesh() model = FakeModel() config = DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=2, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) assert config.cluster_shape == (1, 1) assert config.num_devices == 1 assert config.vocab_size == 8 assert config.allow_force_argmax assert config.max_device_top_k == 32 assert config.sampling_batch_size == 2 assert config.sampling_state_controller is None assert config.sampling_state is None assert config.position_feedback_capable with expect_error(dataclasses.FrozenInstanceError, "cannot assign to field"): config.lane_capacity = 1 with expect_error(TypeError, "DecodeRuntimeConfig"): DecodeRuntime(config=None) def test_config_rejects_inconsistent_collaborators_and_dimensions(expect_error): mesh = FakeMesh() model = FakeModel() model.config.mesh_device = mesh with expect_error(ValueError, "same mesh_device"): DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(FakeMesh()), lane_capacity=2, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) with expect_error(ValueError, "positive integer"): DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=True, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) with expect_error(TypeError, "PageTableLayout"): DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=2, page_table_layout=SimpleNamespace(raw_capacity_width=8, decode_width=8, block_size=32), device_sampling_enabled=True, ) def test_layout_replacement_is_immutable_and_bounded(expect_error): runtime = make_runtime() original = runtime.config replacement = page_table_layout(raw_width=4) runtime.configure_page_table_layout(replacement) assert runtime.config is not original assert original.page_table_layout.raw_capacity_width == 8 assert runtime.config.page_table_layout is replacement with expect_error(ValueError, "block_size"): runtime.configure_page_table_layout(page_table_layout(raw_width=4, block_size=16)) with expect_error(ValueError, "ceiling"): runtime.configure_page_table_layout(page_table_layout(raw_width=9)) with expect_error(ValueError, "decode width"): runtime.configure_page_table_layout(PageTableLayout(32, 4, 8, 1024)) def test_sampling_admission_follows_resolved_configuration(expect_error): runtime = make_runtime(sampling=False) with expect_error(ValueError, "device sampling is disabled"): prepare(runtime, sampling_params=greedy_sampling()) def test_feedback_and_sampling_path_follow_resolved_capabilities(): argmax_runtime = make_runtime() argmax_prepared = prepare(argmax_runtime, sampling_params=greedy_sampling()) assert argmax_prepared.device_feedback assert argmax_prepared.sampling_path == "argmax" model = FakeModel() model.increment_positions = None mesh = FakeMesh() no_feedback = DecodeRuntime( DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=2, page_table_layout=page_table_layout(), device_sampling_enabled=True, force_greedy_top_k=True, ) ) prepared = prepare(no_feedback, sampling_params=greedy_sampling()) assert not prepared.device_feedback assert prepared.sampling_path == "topk" def test_single_and_multi_device_logits_conversion(monkeypatch): single = make_runtime() logits = torch.arange(16, dtype=torch.float32).reshape(1, 1, 2, 8) monkeypatch.setattr(ttnn, "to_torch", lambda value: logits) converted, _ = single._normalize_host_output("single-device", is_tokens=False) assert converted.shape == (2, 1, 8) mesh = SimpleNamespace(shape=(1, 2)) model = FakeModel() model.num_devices = 2 multi = DecodeRuntime( DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=2, page_table_layout=page_table_layout(), device_sampling_enabled=True, ) ) calls = [] monkeypatch.setattr( decode_module, "_concat_host_output", lambda value, shape: calls.append((value, shape)) or logits, ) converted, _ = multi._normalize_host_output("multi-device", is_tokens=False) assert converted.shape == (2, 1, 8) assert calls == [("multi-device", (1, 2))] def test_signatures_expose_ordered_material_and_separate_types(): runtime = make_runtime() prepared = prepare(runtime, sampling_params=greedy_sampling()) program = runtime.program_signature(prepared) trace = runtime.trace_signature(prepared) assert isinstance(program, DecodeProgramSignature) assert isinstance(trace, DecodeTraceSignature) assert program.key_material() == ( ("operation", "decode"), ("batch_size", 2), ("page_table_width", 8), ("sampling_path", "argmax"), ("device_feedback", True), ) assert trace.key_material() == program.key_material() assert runtime.program_signature(prepare(runtime)).sampling_path == "logits" def test_signature_tracks_native_penalty_and_sampled_logprob_program_modes(): runtime = make_runtime() sampling = SamplingParams( temperature=[0.0, 0.0], top_k=[1, 1], top_p=[1.0, 1.0], presence_penalty=[0.5, 0.0], enable_log_probs=[True, False], num_logprobs=[0, -2], ) prepared = prepare(runtime, sampling_params=sampling) native = prepared.prepared_sampling signature = runtime.program_signature(prepared) assert native is not None assert native.logprob_modes == ("sampled_token", "none") assert native.penalties_enabled assert native.log_probs_enabled assert prepared.sampling_path == "topk" assert signature.penalties_enabled assert signature.logprobs_enabled assert signature.key_material()[-2:] == ( ("penalties_enabled", True), ("logprobs_enabled", True), ) def test_configured_topk_policy_is_not_collapsed_to_argmax_by_greedy_temperature(): runtime = make_runtime(force_greedy_top_k=True) sampling = SamplingParams(temperature=[0.0, 0.0], top_k=[32, 32], top_p=[0.08, 0.08]) prepared = prepare(runtime, sampling_params=sampling) assert prepared.sampling_path == "topk" def test_unconfigured_topk_values_use_argmax_for_greedy_temperature(): runtime = make_runtime() sampling = SamplingParams(temperature=[0.0, 0.0], top_k=[32, 32], top_p=[0.08, 0.08]) assert prepare(runtime, sampling_params=sampling).sampling_path == "argmax" def test_sampling_params_are_prepared_once_and_reused_for_kpt(monkeypatch): runtime = make_runtime(force_greedy_top_k=True) calls = [] formatter = decode_module.prepare_sampling_params def formatter_spy(*args, **kwargs): calls.append((args, kwargs)) return formatter(*args, **kwargs) monkeypatch.setattr( decode_module, "prepare_sampling_params", formatter_spy, ) prepared = prepare(runtime, sampling_params=greedy_sampling()) monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh: "mapper") monkeypatch.setattr(ttnn, "from_torch", lambda value, **kwargs: value) runtime._make_host_kpt(prepared) assert len(calls) == 1 def test_tile_padded_sampler_preserves_one_semantic_lane_and_neutralizes_inactive_rows(monkeypatch): mesh = FakeMesh() model = FakeModel() model.config.max_batch_size = 1 model.sampling.config.max_batch_size = 32 runtime = DecodeRuntime( DecodeRuntimeConfig.resolve( model=model, output_reader=OutputReader(mesh), lane_capacity=1, page_table_layout=page_table_layout(), device_sampling_enabled=True, force_greedy_top_k=True, ) ) prepared = runtime.prepare( torch.tensor([11]), torch.tensor([0]), torch.tensor([[3, 4, 5]], dtype=torch.int32), sampling_params=SamplingParams( temperature=0.7, top_k=32, top_p=0.08, presence_penalty=0.25, frequency_penalty=0.5, repetition_penalty=1.2, seed=17, enable_log_probs=True, num_logprobs=1, ), ) sampling = prepared.prepared_sampling assert runtime.config.lane_capacity == 1 assert runtime.config.sampling_batch_size == 32 assert sampling.batch_size == 32 assert sampling.active_rows == 1 assert sampling.active_mask == (True,) + (False,) * 31 assert sampling.top_k == (32,) + (1,) * 31 assert sampling.top_p == pytest.approx((0.08,) + (0.0,) * 31) assert sampling.temperature == pytest.approx((1.0 / 0.7,) + (1.0,) * 31) assert sampling.seeds == (17,) + (None,) * 31 assert sampling.presence_penalty == pytest.approx((0.25,) + (0.0,) * 31) assert sampling.frequency_penalty == pytest.approx((0.5,) + (0.0,) * 31) assert sampling.repetition_penalty == pytest.approx((1.2,) + (1.0,) * 31) assert sampling.enable_log_probs == (True,) + (False,) * 31 assert sampling.num_logprobs == (1,) + (0,) * 31 assert sampling.logprob_modes[1:] == ("none",) * 31 monkeypatch.setattr(ttnn, "ReplicateTensorToMesh", lambda mesh_device: "mapper") monkeypatch.setattr(ttnn, "from_torch", lambda value, **kwargs: value) k, p, temperature = runtime._make_host_kpt(prepared) assert tuple(k.shape) == tuple(p.shape) == tuple(temperature.shape) == (32,) assert k.tolist() == [32] + [1] * 31 assert p.tolist() == pytest.approx([0.08] + [0.0] * 31) assert temperature.tolist() == pytest.approx([1.0 / 0.7] + [1.0] * 31) def test_normalization_preserves_feedback_lookahead_and_inactive_convention(): runtime = make_runtime() page_table = torch.tensor([[10, 11, 99], [20, 21, 98]], dtype=torch.int64) prepared = prepare( runtime, positions=(31, -1), page_table=page_table, sampling_params=greedy_sampling(), ) assert prepared.page_table.dtype == torch.int32 assert prepared.page_table.shape == (2, 8) assert prepared.page_table[0].tolist() == [10, 11, 0, 0, 0, 0, 0, 0] assert prepared.page_table[1].tolist() == [0, 0, 0, 0, 0, 0, 0, 0] def test_normalization_reuses_equal_source_with_same_copy_counts(): runtime = make_runtime() first = prepare(runtime, positions=(0, -1)) second = prepare(runtime, positions=(1, -1)) assert second.page_table is first.page_table def test_normalization_cache_detects_in_place_source_mutation(): runtime = make_runtime() source = torch.tensor([[3, 4], [6, 7]], dtype=torch.int32) first = prepare(runtime, positions=(0, -1), page_table=source) source[0, 0] = 9 second = prepare(runtime, positions=(0, -1), page_table=source) assert second.page_table is not first.page_table assert second.page_table[0, 0].item() == 9 def test_normalization_cache_misses_when_copy_counts_or_feedback_change(): runtime = make_runtime() source = torch.tensor([[3, 4], [6, 7]], dtype=torch.int32) one_block = prepare(runtime, positions=(0, -1), page_table=source) two_blocks = prepare(runtime, positions=(32, -1), page_table=source) no_feedback = prepare(runtime, positions=(31, -1), page_table=source) with_feedback = prepare( runtime, positions=(31, -1), page_table=source, sampling_params=greedy_sampling(), ) assert two_blocks.page_table is not one_block.page_table assert no_feedback.page_table[0, 1].item() == 0 assert with_feedback.page_table is not no_feedback.page_table assert with_feedback.page_table[0, 1].item() == 4 def test_fixed_capacity_and_page_table_capacity_are_validated(expect_error): runtime = make_runtime() with expect_error(ValueError, "must equal lane capacity"): runtime.prepare(torch.tensor([1]), torch.tensor([0]), torch.tensor([[1]])) with expect_error(ValueError, "batches must match"): runtime.prepare( torch.tensor([1, 2]), torch.tensor([0]), torch.tensor([[1], [2]]), ) with expect_error(ValueError, "paged-KV capacity"): prepare(runtime, positions=(8 * 32, -1)) with expect_error(ValueError, "too narrow"): prepare( runtime, positions=(64, -1), page_table=torch.tensor([[1, 2], [0, 0]], dtype=torch.int32), ) def test_preparation_tracks_first_used_page_change_reset_and_ignores_unused_tail(): runtime = make_runtime() first = prepare(runtime, positions=(0, -1), reset=True) assert first.page_table_changed assert first.reset_batch runtime.note_submitted(first) same_semantics = prepare( runtime, positions=(0, -1), page_table=torch.tensor([[3, 90, 91], [88, 87, 86]], dtype=torch.int32), ) assert not same_semantics.page_table_changed assert not same_semantics.reset_batch changed = prepare( runtime, positions=(0, -1), page_table=torch.tensor([[4, 90, 91], [88, 87, 86]], dtype=torch.int32), ) assert changed.page_table_changed def test_submission_state_tracks_last_table_despite_stale_prepare_change_hint(): runtime = make_runtime() baseline = prepare( runtime, positions=(0, -1), page_table=torch.tensor([[3], [0]], dtype=torch.int32), ) runtime.note_submitted(baseline) changed = prepare( runtime, positions=(0, -1), page_table=torch.tensor([[4], [0]], dtype=torch.int32), ) back_to_baseline = prepare( runtime, positions=(0, -1), page_table=torch.tensor([[3], [0]], dtype=torch.int32), ) assert changed.page_table_changed assert not back_to_baseline.page_table_changed runtime.note_submitted(changed) runtime.note_submitted(back_to_baseline) assert not prepare(runtime, page_table=back_to_baseline.page_table).page_table_changed assert prepare( runtime, positions=(0, -1), page_table=torch.tensor([[4], [0]], dtype=torch.int32), ).page_table_changed def test_capture_plan_describes_full_step_refresh_and_typed_persistent_inputs(monkeypatch): runtime = make_runtime() prepared = prepare(runtime, sampling_params=greedy_sampling()) device = DecodeDeviceInputs("tokens", "positions", "rotary", "page_table") monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda request: "host") monkeypatch.setattr(runtime, "_stage_inputs_and_kpt", lambda host, request: (device, "kpt")) def run_body( inputs, prepared, kpt, *, device_feedback, count_tokens=True, advance_seeds=True, ): return "captured" monkeypatch.setattr(runtime, "_run_body", run_body) plan = runtime.capture_plan(prepared) persistent = plan.prepare_inputs() assert persistent.device_inputs is device assert persistent.kpt == "kpt" sampling = prepared.prepared_sampling assert sampling is not None assert persistent.kpt_signature == [(sampling.top_k, sampling.top_p, sampling.temperature)] assert plan.capture(persistent) == "captured" assert plan.refresh_policy.every_replay == ("sampling",) assert plan.refresh_policy.full_on_batch_reset assert plan.refresh_policy.full_on_graph_switch assert plan.refresh_policy.full_without_device_feedback assert plan.refresh_policy.refresh_page_table_on_change def test_trace_refresh_skips_unchanged_sampling_values(monkeypatch): runtime = make_runtime(force_greedy_top_k=True) prepared = prepare(runtime, sampling_params=greedy_sampling()) sampling = prepared.prepared_sampling assert sampling is not None persistent = DecodePersistentInputs( device_inputs=DecodeDeviceInputs("tokens", "positions", "rotary", "page_table"), kpt="kpt", kpt_signature=[(sampling.top_k, sampling.top_p, sampling.temperature)], ) def fail_refresh_kpt(device_kpt, prepared): pytest.fail("unchanged KPT was refreshed") monkeypatch.setattr(runtime, "_refresh_kpt", fail_refresh_kpt) runtime.refresh_trace( persistent, prepared, SimpleNamespace(full=False, page_table=False), ) def test_eager_invoke_returns_owned_result_and_advances_submission_state(monkeypatch): runtime = make_runtime() prepared = prepare(runtime) device = DecodeDeviceInputs("tokens", "positions", "rotary", "page_table") calls = [] monkeypatch.setattr(runtime, "_prepare_inputs_host", lambda request: "host") monkeypatch.setattr(runtime, "_stage_inputs_and_kpt", lambda host, request: (device, None)) monkeypatch.setattr( runtime, "_run_body", lambda inputs, prepared, kpt, *, device_feedback, **kwargs: calls.append(device_feedback) or ("raw", None), ) result = runtime.invoke(prepared) assert isinstance(result, InvocationResult) assert result.value == ("raw", None) assert result.owned == (("raw", None), (device, None)) assert not result.is_tokens assert calls == [False] assert not prepare(runtime, page_table=prepared.page_table).page_table_changed def test_blocking_consume_normalizes_logits_and_releases_owned_values(monkeypatch): runtime = make_runtime() host_logits = torch.arange(16, dtype=torch.float32).reshape(1, 1, 2, 8) released = [] monkeypatch.setattr(runtime.config.output_reader, "read", lambda value, *, blocking: (host_logits, "probs")) monkeypatch.setattr(runtime, "_release_or_retain_transient", lambda value: released.append(value) or []) result = InvocationResult(value="raw", owned="owned", is_tokens=False) logits, log_probs = runtime.consume(result) assert logits.shape == (2, 1, 8) assert log_probs == "probs" assert released == ["owned"] def test_sampled_token_logprobs_are_flattened_to_lane_row_order(): runtime = make_runtime() host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32) host_log_probs = torch.tensor([[[[-0.25, -0.75]]]], dtype=torch.bfloat16) tokens, log_probs = runtime._normalize_host_output( (host_tokens, host_log_probs), is_tokens=True, ) assert tokens.tolist() == [7, 8] assert tokens.dtype == torch.int64 assert log_probs.tolist() == pytest.approx([-0.25, -0.75]) assert log_probs.dtype == torch.float32 def test_raw_blocking_and_async_leases_release_exact_records(monkeypatch): runtime = make_runtime() deallocated = [] monkeypatch.setattr( decode_module, "best_effort_deallocate_owned_tensors", lambda values, completed: deallocated.append(values) or [], ) first = InvocationResult(value=object(), owned="first-owned", is_tokens=False) assert runtime.consume(first, read_from_device=False) is first.value monkeypatch.setattr(runtime.config.output_reader, "read", lambda value, *, blocking: "first-host") assert runtime.read_decode_output(first.value) == "first-host" second = InvocationResult(value=object(), owned="second-owned", is_tokens=True) runtime.consume(second, read_from_device=False) host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32) pending = PendingRead(value=(host_tokens, None), events=("event",), sequence=4, _owner=object()) monkeypatch.setattr(runtime.config.output_reader, "submit", lambda value: pending) monkeypatch.setattr(runtime.config.output_reader, "complete", lambda value: pending.value) host, events = runtime.read_decode_output(second.value, async_read=True) assert host is pending.value assert events == ["event"] tokens, log_probs = runtime.process_decode_output_host(host, is_tokens=True) assert tokens.tolist() == [7, 8] assert tokens.dtype == torch.int64 assert log_probs is None assert deallocated == [ (first.value, "first-owned"), (second.value, "second-owned"), ] def test_async_trace_lease_never_releases_borrowed_trace_output(monkeypatch): runtime = make_runtime() deallocated = [] raw = object() host_tokens = torch.tensor([[[[7], [8]]]], dtype=torch.int32) pending = PendingRead(value=(host_tokens, None), events=("event",), sequence=4, _owner=object()) monkeypatch.setattr( decode_module, "best_effort_deallocate_owned_tensors", lambda values, completed: deallocated.append(values) or [], ) monkeypatch.setattr(runtime.config.output_reader, "submit", lambda value: pending) monkeypatch.setattr(runtime.config.output_reader, "complete", lambda value: pending.value) result = InvocationResult(value=raw, owned=None, is_tokens=True) assert runtime.consume(result, read_from_device=False) is raw host, events = runtime.read_decode_output(raw, async_read=True) assert events == ["event"] tokens, log_probs = runtime.process_decode_output_host(host, is_tokens=True) assert tokens.tolist() == [7, 8] assert log_probs is None assert deallocated == [] def test_failed_transient_release_blocks_use_and_cleanup_retries(monkeypatch, expect_error): runtime = make_runtime() class FakeTensor: pass tensor = FakeTensor() attempts = [] monkeypatch.setattr(decode_module.ttnn, "Tensor", FakeTensor) def deallocate(value): attempts.append(value) if len(attempts) == 1: raise RuntimeError("release failed") monkeypatch.setattr(decode_module.ttnn, "deallocate", deallocate) failures = runtime._release_or_retain_transient(tensor) assert [str(error) for error in failures] == ["release failed"] assert runtime.transient_orphan_count == 1 with expect_error(RuntimeError, "unreleased transient"): prepare(runtime) runtime.cleanup_transients() assert attempts == [tensor, tensor] assert runtime.transient_orphan_count == 0 assert prepare(runtime).sampling_path == "logits"