# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC # SPDX-License-Identifier: Apache-2.0 """Generic host-only tensor-model contract for migrated sibling models.""" import inspect from dataclasses import dataclass, fields from types import SimpleNamespace from unittest.mock import MagicMock, call import pytest from models.common.models.deepseek_r1_distill_qwen_14b import model as deepseek_model from models.common.models.llama32_1b import model as llama32_model from models.common.models.llama32_3b import model as llama32_3b_model from models.common.models.llama33_70b import model as llama33_70b_model from models.common.models.mistral_7b import model as mistral_model from models.common.models.phi4 import model as phi4_model from models.common.models.qwen2_7b import model as qwen2_model from models.common.models.qwen3_32b import model as qwen3_32b_model from models.common.models.qwen25_7b import model as qwen25_model from models.common.models.qwen25_72b import model as qwen25_72b_model from models.common.models.qwen25_coder_32b import model as qwen25_coder_32b_model MODEL_CONTRACTS = { "llama32_1b": SimpleNamespace( module=llama32_model, model_class=llama32_model.Llama32_1BTransformer1D, config_class=llama32_model.Llama32_1BTransformer1DConfig, attention_config_class=llama32_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_llama32_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "llama32_3b": SimpleNamespace( module=llama32_3b_model, model_class=llama32_3b_model.Llama32_3BTransformer1D, config_class=llama32_3b_model.Llama32_3BTransformer1DConfig, attention_config_class=llama32_3b_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_llama32_config(module=llama32_3b_model, **kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch, module=llama32_3b_model), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "llama33_70b": SimpleNamespace( module=llama33_70b_model, model_class=llama33_70b_model.Llama33_70BTransformer1D, config_class=llama33_70b_model.Llama33_70BTransformer1DConfig, attention_config_class=llama33_70b_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_llama32_config(module=llama33_70b_model, **kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_llama32_model(monkeypatch, module=llama33_70b_model), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "qwen2_7b": SimpleNamespace( module=qwen2_model, model_class=qwen2_model.Qwen2_7B, config_class=qwen2_model.Qwen2_7BTransformerConfig, attention_config_class=qwen2_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen2_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "qwen25_7b": SimpleNamespace( module=qwen25_model, model_class=qwen25_model.Qwen25_7B, config_class=qwen25_model.Qwen25_7BTransformerConfig, attention_config_class=qwen25_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen2_config(module=qwen25_model, **kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch, module=qwen25_model), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "qwen25_72b": SimpleNamespace( module=qwen25_72b_model, model_class=qwen25_72b_model.Qwen25_72B, config_class=qwen25_72b_model.Qwen25_72BConfig, attention_config_class=qwen25_72b_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen25_72b_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen25_72b_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "qwen25_coder_32b": SimpleNamespace( module=qwen25_coder_32b_model, model_class=qwen25_coder_32b_model.Qwen25Coder32B, config_class=qwen25_coder_32b_model.Qwen25Coder32BConfig, attention_config_class=qwen25_coder_32b_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen25_coder_32b_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen25_coder_32b_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "qwen3_32b": SimpleNamespace( module=qwen3_32b_model, model_class=qwen3_32b_model.Qwen3_32B, config_class=qwen3_32b_model.Qwen3_32BConfig, attention_config_class=qwen3_32b_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen3_32b_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen3_32b_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "deepseek_r1_distill_qwen_14b": SimpleNamespace( module=deepseek_model, model_class=deepseek_model.DeepSeekR1Qwen14B, config_class=deepseek_model.DeepSeekR1Qwen14BTransformerConfig, attention_config_class=deepseek_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_qwen2_config(module=deepseek_model, **kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_qwen2_model(monkeypatch, module=deepseek_model), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "mistral_7b": SimpleNamespace( module=mistral_model, model_class=mistral_model.Mistral7B, config_class=mistral_model.Mistral7BTransformerConfig, attention_config_class=mistral_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_mistral_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_mistral_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), "phi4": SimpleNamespace( module=phi4_model, model_class=phi4_model.Phi4Transformer, config_class=phi4_model.Phi4TransformerConfig, attention_config_class=phi4_model.Attention1DConfig, make_attention_config=lambda **kwargs: _make_llama32_attention_config(**kwargs), make_config=lambda **kwargs: _make_phi4_config(**kwargs), make_layer=lambda attention_config=None: _make_llama32_layer(attention_config), construct_model=lambda monkeypatch: _construct_phi4_model(monkeypatch), expected_module_names=( "layer[0].attn_norm", "layer[0].attention", "layer[0].ff_norm", "layer[0].mlp", "layer[1].attn_norm", "layer[1].attention", "layer[1].ff_norm", "layer[1].mlp", "final_norm", "lm_head", ), ), } @pytest.fixture(params=MODEL_CONTRACTS.items(), ids=lambda item: item[0]) def contract(request): return request.param[1] @dataclass(frozen=True) class _PagedAttentionConfig: block_size: int max_num_blocks: int def _make_llama32_attention_config(*, n_kv_heads=8, kv_cache=None): return SimpleNamespace( n_kv_heads=n_kv_heads, use_vllm_paged_kv_cache=True, paged_attention_config=_PagedAttentionConfig(block_size=32, max_num_blocks=128), kv_cache=kv_cache, ) def _make_llama32_layer(attention_config=None): attention_config = attention_config or _make_llama32_attention_config() attention = SimpleNamespace(config=attention_config, kv_cache=attention_config.kv_cache) return SimpleNamespace( attention_norm=object(), attention=attention, self_attn=attention, ff_norm=object(), feed_forward=object(), ) def _make_llama32_config(*, module=llama32_model, n_layers=1, num_devices=2, n_kv_heads=8, sampling_config=None): block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] config_class = next( getattr(module, name) for name in ( "Llama32_1BTransformer1DConfig", "Llama32_3BTransformer1DConfig", "Llama33_70BTransformer1DConfig", ) if hasattr(module, name) ) return config_class( n_layers=n_layers, vocab_size=128256, max_batch_size=4, max_seq_len=4096, dim=2048, num_devices=num_devices, mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices), embedding_config=object(), rope_config=object(), block_configs=block_configs, norm_config=object(), lm_head_config=object(), sampling_config=sampling_config, ) def _construct_llama32_model(monkeypatch, *, module=llama32_model): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } for owner_name, sentinel_name in ( ("Embedding1D", "embedding"), ("RotarySetup1D", "rope_setup"), ("TransformerBlock1D", "layer"), ("RMSNorm1D", "norm"), ("LMHead1D", "lm_head"), ("Sampling1D", "sampling"), ): monkeypatch.setattr( getattr(module, owner_name), "from_config", MagicMock(return_value=sentinels[sentinel_name]), ) config = _make_llama32_config(module=module, num_devices=1, sampling_config=object()) model_class = next( getattr(module, name) for name in ( "Llama32_1BTransformer1D", "Llama32_3BTransformer1D", "Llama33_70BTransformer1D", ) if hasattr(module, name) ) return model_class(config), config, sentinels def _make_qwen2_config(*, module=qwen2_model, n_layers=1, num_devices=2, n_kv_heads=4, sampling_config=None): block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] config_class = ( getattr(module, "Qwen2_7BTransformerConfig", None) or getattr(module, "Qwen25_7BTransformerConfig", None) or module.DeepSeekR1Qwen14BTransformerConfig ) return config_class( n_layers=n_layers, vocab_size=152064, max_batch_size=4, max_seq_len=4096, dim=3584, num_devices=num_devices, mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices), embedding_config=object(), rope_config=object(), block_configs=block_configs, norm_config=object(), lm_head_config=object(), sampling_config=sampling_config, tt_ccl=object(), ) def _construct_qwen2_model(monkeypatch, *, module=qwen2_model): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } layer_class = ( "Qwen2_7BDecoderLayer" if module is qwen2_model else ("Qwen25_7BDecoderLayer" if module is qwen25_model else "DeepSeekR1Qwen14BDecoderLayer") ) for owner_name, sentinel_name in ( ("Embedding1D", "embedding"), ("RotarySetup1D", "rope_setup"), (layer_class, "layer"), ("RMSNorm1D", "norm"), ("LMHead1D", "lm_head"), ("Sampling1D", "sampling"), ): monkeypatch.setattr( getattr(module, owner_name), "from_config", MagicMock(return_value=sentinels[sentinel_name]), ) config = _make_qwen2_config(module=module, sampling_config=object()) model_class = getattr(module, "Qwen2_7B", None) or getattr(module, "Qwen25_7B", None) or module.DeepSeekR1Qwen14B return model_class(config), config, sentinels def _make_qwen25_72b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None): del sampling_config block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] return qwen25_72b_model.Qwen25_72BConfig( hf_model_id="Qwen/Qwen2.5-72B-Instruct", dim=8192, n_heads=64, n_kv_heads=n_kv_heads, head_dim=128, hidden_dim=29568, vocab_size=152064, rms_norm_eps=1e-6, rope_theta=1_000_000.0, num_hidden_layers=n_layers, max_batch_size=4, max_seq_len=4096, rope_table_len=8192, num_devices=num_devices, mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices), block_configs=block_configs, ) def _construct_qwen25_72b_model(monkeypatch): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } monkeypatch.setattr(qwen25_72b_model, "get_tt_ccl", MagicMock(return_value=object())) monkeypatch.setattr(qwen25_72b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"])) config = _make_qwen25_72b_config() return ( qwen25_72b_model.Qwen25_72B( config, sentinels["embedding"], sentinels["rope_setup"], [sentinels["layer"]], sentinels["norm"], sentinels["lm_head"], config.mesh_device, ), config, sentinels, ) def _make_qwen25_coder_32b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None): del sampling_config block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] return qwen25_coder_32b_model.Qwen25Coder32BConfig( hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct", dim=5120, n_heads=40, n_kv_heads=n_kv_heads, head_dim=128, hidden_dim=27648, vocab_size=152064, rms_norm_eps=1e-6, rope_theta=1_000_000.0, num_hidden_layers=n_layers, max_batch_size=4, max_seq_len=4096, rope_table_len=8192, num_devices=num_devices, mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices), block_configs=block_configs, ) def _construct_qwen25_coder_32b_model(monkeypatch): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } monkeypatch.setattr(qwen25_coder_32b_model, "get_tt_ccl", MagicMock(return_value=object())) monkeypatch.setattr(qwen25_coder_32b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"])) config = _make_qwen25_coder_32b_config() return ( qwen25_coder_32b_model.Qwen25Coder32B( config, sentinels["embedding"], sentinels["rope_setup"], [sentinels["layer"]], sentinels["norm"], sentinels["lm_head"], config.mesh_device, ), config, sentinels, ) def _make_qwen3_32b_config(*, n_layers=1, num_devices=8, n_kv_heads=8, sampling_config=None): del sampling_config block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] return qwen3_32b_model.Qwen3_32BConfig( hf_model_id="Qwen/Qwen3-32B", dim=5120, n_heads=64, n_kv_heads=n_kv_heads, head_dim=128, hidden_dim=27648, vocab_size=151936, rms_norm_eps=1e-6, rope_theta=1_000_000.0, num_hidden_layers=n_layers, max_batch_size=4, max_seq_len=4096, rope_table_len=8192, num_devices=num_devices, mesh_device=SimpleNamespace(get_num_devices=lambda: num_devices), block_configs=block_configs, ) def _construct_qwen3_32b_model(monkeypatch): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } monkeypatch.setattr(qwen3_32b_model, "get_tt_ccl", MagicMock(return_value=object())) monkeypatch.setattr(qwen3_32b_model, "Sampling1D", MagicMock(return_value=sentinels["sampling"])) monkeypatch.setattr(qwen3_32b_model.weight_utils, "lm_head_padded_vocab_size", MagicMock(return_value=152064)) config = _make_qwen3_32b_config() return ( qwen3_32b_model.Qwen3_32B( config, sentinels["embedding"], sentinels["rope_setup"], [sentinels["layer"]], sentinels["norm"], sentinels["lm_head"], config.mesh_device, ), config, sentinels, ) def _make_mistral_config(*, n_layers=1, num_devices=2, n_kv_heads=8, sampling_config=None): block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] return mistral_model.Mistral7BTransformerConfig( n_layers=n_layers, vocab_size=32768, max_batch_size=4, max_seq_len=4096, dim=4096, num_devices=num_devices, mesh_device=object(), embedding_config=object(), rope_config=object(), block_configs=block_configs, norm_config=object(), lm_head_config=object(), sampling_config=sampling_config, tt_ccl=object(), ) def _construct_mistral_model(monkeypatch): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } for owner_name, sentinel_name in ( ("Embedding1D", "embedding"), ("RotarySetup1D", "rope_setup"), ("Mistral7BDecoderLayer", "layer"), ("RMSNorm1D", "norm"), ("LMHead1D", "lm_head"), ("Sampling1D", "sampling"), ): monkeypatch.setattr( getattr(mistral_model, owner_name), "from_config", MagicMock(return_value=sentinels[sentinel_name]), ) config = _make_mistral_config(sampling_config=object()) return mistral_model.Mistral7B(config), config, sentinels def _make_phi4_config(*, n_layers=1, num_devices=2, n_kv_heads=10, sampling_config=None): block_configs = [ SimpleNamespace(attention_config=_make_llama32_attention_config(n_kv_heads=n_kv_heads)) for _ in range(n_layers) ] return phi4_model.Phi4TransformerConfig( n_layers=n_layers, vocab_size=100352, max_batch_size=4, max_seq_len=4096, dim=5120, num_devices=num_devices, mesh_device=object(), embedding_config=object(), rope_config=object(), block_configs=block_configs, norm_config=object(), lm_head_config=object(), sampling_config=sampling_config, tt_ccl=object(), ) def _construct_phi4_model(monkeypatch): sentinels = { "embedding": object(), "rope_setup": object(), "layer": _make_llama32_layer(), "norm": object(), "lm_head": object(), "sampling": object(), } for owner_name, sentinel_name in ( ("Embedding1D", "embedding"), ("RotarySetup1D", "rope_setup"), ("Phi4DecoderLayer", "layer"), ("RMSNorm1D", "norm"), ("LMHead1D", "lm_head"), ("Sampling1D", "sampling"), ): monkeypatch.setattr( getattr(phi4_model, owner_name), "from_config", MagicMock(return_value=sentinels[sentinel_name]), ) config = _make_phi4_config(sampling_config=object()) return phi4_model.Phi4Transformer(config), config, sentinels def test_config_exposes_complete_runtime_metadata(contract): names = {field.name for field in fields(contract.config_class)} assert { "dim", "mesh_device", "num_devices", "n_layers", "max_batch_size", "max_seq_len", "block_configs", } <= names assert "n_kv_heads" in {field.name for field in fields(contract.attention_config_class)} config = contract.make_config(n_layers=2, num_devices=2, n_kv_heads=8) assert len(config.block_configs) == config.n_layers for block in config.block_configs: assert hasattr(block.attention_config, "n_kv_heads") assert block.attention_config.n_kv_heads % config.num_devices == 0 @pytest.mark.parametrize( "method,names,keyword_only,defaults", [ ( "prefill_forward", ( "self", "x_embed", "rot_mats", "user_id", "page_table", "chunk_page_table", "chunk_start_idx", "get_last_token", "batch_size", "chunk_start_idx_tensor", "last_token_slice", "last_token_index", ), (), { "user_id": 0, "page_table": None, "chunk_page_table": None, "chunk_start_idx": None, "get_last_token": -1, "batch_size": 1, "chunk_start_idx_tensor": None, "last_token_slice": None, "last_token_index": None, }, ), ( "post_process_prefill_output", ("self", "hidden_states", "last_token_idx", "last_token_slice", "last_token_index"), (), {"last_token_slice": None, "last_token_index": None}, ), ( "post_process_batched_prefill_output", ( "self", "hidden_states", "last_token_idx_list", "padded_batch", "prefill_seq_len", "last_token_slice", "last_token_index", ), (), {"last_token_slice": None, "last_token_index": None}, ), ("set_kv_cache", ("self", "kv_cache"), (), {}), ( "configure_paged_attention", ("self", "block_size", "max_num_blocks"), ("block_size", "max_num_blocks"), {}, ), ("prepare_prefill_rot_mats", ("self", "position_indices"), (), {}), ("iter_executor_named_modules", ("self",), (), {}), ], ) def test_method_signatures_are_exact(contract, method, names, keyword_only, defaults): signature = inspect.signature(getattr(contract.model_class, method)) parameters = signature.parameters assert tuple(parameters) == names for name, parameter in parameters.items(): expected_kind = ( inspect.Parameter.KEYWORD_ONLY if name in keyword_only else inspect.Parameter.POSITIONAL_OR_KEYWORD ) assert parameter.kind is expected_kind expected_default = defaults.get(name, inspect.Parameter.empty) assert parameter.default == expected_default annotations = { "prefill_forward": ( ( inspect.Parameter.empty, "ttnn.Tensor", "tuple[ttnn.Tensor, ttnn.Tensor]", "int", "ttnn.Tensor | None", "ttnn.Tensor | None", "int | None", "int", "int", "ttnn.Tensor | None", "tuple[ttnn.Tensor, ttnn.Tensor] | None", "ttnn.Tensor | None", ), "ttnn.Tensor", ), "post_process_prefill_output": ( ( inspect.Parameter.empty, "ttnn.Tensor", "int", "tuple[ttnn.Tensor, ttnn.Tensor] | None", "ttnn.Tensor | None", ), "ttnn.Tensor", ), "post_process_batched_prefill_output": ( ( inspect.Parameter.empty, "ttnn.Tensor", "list[int]", "int", "int", "tuple[ttnn.Tensor, ttnn.Tensor] | None", "ttnn.Tensor | None", ), "ttnn.Tensor", ), "set_kv_cache": ((inspect.Parameter.empty, "list | None"), "None"), "configure_paged_attention": ((inspect.Parameter.empty, "int", "int"), "None"), "prepare_prefill_rot_mats": ( (inspect.Parameter.empty, "ttnn.Tensor"), "tuple[ttnn.Tensor, ttnn.Tensor]", ), "iter_executor_named_modules": ( (inspect.Parameter.empty,), inspect.Signature.empty, ), } expected_parameter_annotations, expected_return_annotation = annotations[method] assert tuple(parameter.annotation for parameter in parameters.values()) == expected_parameter_annotations assert signature.return_annotation == expected_return_annotation def test_required_methods_exist(contract): for method in ( "set_kv_cache", "configure_paged_attention", "prepare_prefill_rot_mats", "iter_executor_named_modules", "embed_decode", "embed_prefill", "gather_and_untilize_logits", "increment_positions", ): assert callable(getattr(contract.model_class, method)) def test_constructed_model_resolves_runtime_surface(contract, monkeypatch): model, config, sentinels = contract.construct_model(monkeypatch) assert model.config is config assert model.embedding is sentinels["embedding"] assert model.rope_setup is sentinels["rope_setup"] assert model.layers == [sentinels["layer"]] assert model.norm is sentinels["norm"] assert model.lm_head is sentinels["lm_head"] assert model.sampling is sentinels["sampling"] assert model.supports_on_device_sampling assert model.mesh_device is config.mesh_device assert model.vocab_size == config.vocab_size assert model.n_layers == config.n_layers assert model.num_devices == config.num_devices assert model.model_args is None def test_named_modules_are_complete_unique_and_ordered(contract): layers = [contract.make_layer(), contract.make_layer()] model = SimpleNamespace(layers=layers, norm=object(), lm_head=object()) named = list(contract.model_class.iter_executor_named_modules(model)) assert tuple(name for name, _ in named) == contract.expected_module_names assert len(named) == 4 * len(layers) + 2 assert len({name for name, _ in named}) == len(named) assert tuple(module for _, module in named) == ( layers[0].attention_norm, layers[0].attention, layers[0].ff_norm, layers[0].feed_forward, layers[1].attention_norm, layers[1].attention, layers[1].ff_norm, layers[1].feed_forward, model.norm, model.lm_head, ) def test_named_modules_without_layers_yields_nothing(contract): assert list(contract.model_class.iter_executor_named_modules(SimpleNamespace())) == [] def test_set_kv_cache_binds_identity_and_unbinds_idempotently(contract): layers = [contract.make_layer(), contract.make_layer()] model = SimpleNamespace(layers=layers) cache = [[object(), object()], [object(), object()]] contract.model_class.set_kv_cache(model, cache) for layer, expected in zip(layers, cache): bound = layer.attention.config.kv_cache assert bound == tuple(expected) assert bound[0] is expected[0] assert bound[1] is expected[1] assert layer.attention.kv_cache is bound contract.model_class.set_kv_cache(model, None) contract.model_class.set_kv_cache(model, None) assert all(layer.attention.config.kv_cache is None for layer in layers) assert all(layer.attention.kv_cache is None for layer in layers) def test_set_kv_cache_rejects_wrong_layer_count_before_binding(contract, expect_error): layers = [contract.make_layer(), contract.make_layer()] model = SimpleNamespace(layers=layers) with expect_error(ValueError, "model has 2 layers"): contract.model_class.set_kv_cache(model, [[object(), object()]]) assert all(layer.attention.config.kv_cache is None for layer in layers) assert all(layer.attention.kv_cache is None for layer in layers) @pytest.mark.parametrize("bad_pair", [[object()], object()]) def test_set_kv_cache_validates_all_pairs_before_binding(contract, bad_pair, expect_error): layers = [contract.make_layer(), contract.make_layer()] model = SimpleNamespace(layers=layers) with expect_error((TypeError, ValueError), "layer 1.*K/V tensor"): contract.model_class.set_kv_cache(model, [[object(), object()], bad_pair]) assert all(layer.attention.config.kv_cache is None for layer in layers) assert all(layer.attention.kv_cache is None for layer in layers) def test_configure_paged_attention_updates_construction_and_live_configs_and_rejects_bound_cache( contract, expect_error ): construction = contract.make_attention_config() live = contract.make_attention_config() model = SimpleNamespace( config=SimpleNamespace(block_configs=(SimpleNamespace(attention_config=construction),)), layers=(contract.make_layer(live),), ) contract.model_class.configure_paged_attention(model, block_size=16, max_num_blocks=200) assert construction.paged_attention_config.block_size == live.paged_attention_config.block_size == 16 assert construction.paged_attention_config.max_num_blocks == live.paged_attention_config.max_num_blocks == 200 bound = object() live.kv_cache = (bound, bound) with expect_error(RuntimeError, "already has a bound KV cache"): contract.model_class.configure_paged_attention(model, block_size=32, max_num_blocks=128) assert construction.paged_attention_config.block_size == live.paged_attention_config.block_size == 16 def test_prepare_prefill_rot_mats_gathers_device_rows(contract, monkeypatch): position_indices = object() cos_matrix, sin_matrix = object(), object() cos_rows, sin_rows = object(), object() cos_4d, sin_4d = object(), object() fake_ttnn = SimpleNamespace( TILE_LAYOUT=object(), embedding=MagicMock(side_effect=[cos_rows, sin_rows]), unsqueeze_to_4D=MagicMock(side_effect=[cos_4d, sin_4d]), ) rope = SimpleNamespace(cos_matrix=cos_matrix, sin_matrix=sin_matrix, load_device_weights=MagicMock()) monkeypatch.setattr(contract.module, "ttnn", fake_ttnn) result = contract.model_class.prepare_prefill_rot_mats(SimpleNamespace(rope_setup=rope), position_indices) rope.load_device_weights.assert_called_once_with() assert fake_ttnn.embedding.call_args_list == [ call(position_indices, cos_matrix, layout=fake_ttnn.TILE_LAYOUT), call(position_indices, sin_matrix, layout=fake_ttnn.TILE_LAYOUT), ] assert result == (cos_4d, sin_4d)