| |
| |
|
|
| from dataclasses import FrozenInstanceError, fields |
|
|
| import pytest |
|
|
| import ttnn |
| from models.common.llm_runtime import config as runtime_config |
| from models.common.llm_runtime.config import PagedKVCacheConfig, PageTableLayout, TraceConfig, WarmupConfig |
| from models.common.models.llama3_8b.executor import Llama3ExecutorConfig |
|
|
|
|
| class _TraceConfigSubclass(TraceConfig): |
| pass |
|
|
|
|
| def _paged_config(**overrides): |
| kwargs = { |
| "block_size": 32, |
| "max_num_blocks": 1024, |
| "dtype": ttnn.bfloat8_b, |
| } |
| kwargs.update(overrides) |
| return PagedKVCacheConfig(**kwargs) |
|
|
|
|
| def test_executor_config_has_exact_static_policy_owners_and_is_frozen(expect_error): |
| config = Llama3ExecutorConfig( |
| trace=TraceConfig(mode="all"), |
| warmup=WarmupConfig(), |
| paged_kv_cache=_paged_config(), |
| device_sampling_enabled=True, |
| ) |
|
|
| assert [field.name for field in fields(config)] == [ |
| "trace", |
| "warmup", |
| "paged_kv_cache", |
| "device_sampling_enabled", |
| "allow_batched_prefill_with_device_sampling_for_diagnostics", |
| ] |
| assert not config.allow_batched_prefill_with_device_sampling_for_diagnostics |
| forbidden = { |
| "model", |
| "mesh_device", |
| "hf_model", |
| "tokenizer", |
| "dtype", |
| "n_layers", |
| "sampling_config", |
| "sampling_output_dtype", |
| } |
| assert forbidden.isdisjoint(field.name for field in fields(config)) |
| assert not hasattr(runtime_config, "LLMGraphCompilerConfig") |
| assert not hasattr(runtime_config, "LLMExecutorConfig") |
| assert not hasattr(runtime_config, "Sampling1DConfig") |
| with expect_error(FrozenInstanceError, ""): |
| config.device_sampling_enabled = False |
|
|
|
|
| @pytest.mark.parametrize( |
| ("field_name", "invalid_value"), |
| [ |
| ("trace", WarmupConfig()), |
| ("trace", _TraceConfigSubclass()), |
| ("warmup", TraceConfig()), |
| ("paged_kv_cache", WarmupConfig()), |
| ], |
| ) |
| def test_executor_config_rejects_non_exact_nested_config_types(field_name, invalid_value, expect_error): |
| values = { |
| "trace": TraceConfig(), |
| "warmup": WarmupConfig(), |
| "paged_kv_cache": _paged_config(), |
| "device_sampling_enabled": False, |
| } |
| values[field_name] = invalid_value |
|
|
| with expect_error(TypeError, rf"{field_name} must be exactly"): |
| Llama3ExecutorConfig(**values) |
|
|
|
|
| @pytest.mark.parametrize( |
| ("mode", "prefill", "decode"), |
| [("none", False, False), ("decode_only", False, True), ("all", True, True)], |
| ) |
| def test_trace_config_selects_static_coverage(mode, prefill, decode, expect_error): |
| config = TraceConfig(mode=mode) |
|
|
| assert config.prefill_enabled is prefill |
| assert config.decode_enabled is decode |
| with expect_error(FrozenInstanceError, ""): |
| config.mode = "none" |
|
|
|
|
| def test_trace_config_rejects_unknown_mode(expect_error): |
| with expect_error(ValueError, "Unsupported trace mode"): |
| TraceConfig(mode="prefill_only") |
|
|
|
|
| def test_warmup_config_keeps_model_derived_defaults_and_is_deeply_immutable(expect_error): |
| config = WarmupConfig() |
|
|
| assert config.prefill_seq_lens is None |
| assert config.prefill_batch_sizes == (1, 2, 4, 8, 16, 32) |
| assert config.include_decode_top_k is False |
| with expect_error(TypeError, "must be a tuple"): |
| WarmupConfig(prefill_batch_sizes=[1, 2]) |
|
|
|
|
| def test_paged_kv_config_has_plan_fields_and_resolved_capacity(expect_error): |
| unresolved = _paged_config() |
| resolved = _paged_config(num_blocks=512) |
|
|
| assert [field.name for field in fields(unresolved)] == [ |
| "block_size", |
| "max_num_blocks", |
| "dtype", |
| "memory_config", |
| "num_blocks", |
| ] |
| assert unresolved.memory_config == ttnn.DRAM_MEMORY_CONFIG |
| assert not unresolved.is_resolved() |
| assert resolved.is_resolved() |
| with expect_error(FrozenInstanceError, ""): |
| resolved.num_blocks = 256 |
|
|
|
|
| def test_paged_kv_config_rejects_invalid_capacity(expect_error): |
| with expect_error(ValueError, "exceeds max_num_blocks"): |
| _paged_config(num_blocks=1025) |
| with expect_error(ValueError, "block_size"): |
| _paged_config(block_size=0) |
|
|
|
|
| def test_page_table_layout_is_resolved_without_warmup_policy(): |
| layout = PageTableLayout.resolve( |
| block_size=32, |
| model_max_sequence_length=4096, |
| physical_num_blocks=100, |
| max_prefill_chunk_size=2048, |
| ) |
|
|
| assert layout.raw_capacity_width == 100 |
| assert layout.decode_width == 104 |
| assert layout.prefill_width == 168 |
|
|
