# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import os import pytest import torch from loguru import logger import ttnn from models.common.utility_functions import comp_allclose, comp_pcc from models.tt_transformers.tests.test_utils import get_ref_model_dype from models.tt_transformers.tt.attention import Attention from models.tt_transformers.tt.ccl import TT_CCL from models.tt_transformers.tt.common import Mode, PagedAttentionConfig, precompute_freqs from models.tt_transformers.tt.model_config import ModelArgs from models.tt_transformers.tt.prefetcher import Prefetcher from models.tt_transformers.tt.rope import HfRotarySetup, RotarySetup @torch.no_grad() @pytest.mark.parametrize( "use_prefetcher", ([False]), ) @pytest.mark.parametrize( "mesh_device", [ {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get( os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids()) ) ], indirect=True, ) @pytest.mark.parametrize( "paged_attention", ( True, False, ), ids=( "paged_attention", "default_attention", ), ) @pytest.mark.parametrize( "page_params", [{"page_block_size": 32, "page_max_num_blocks": 1024}], ) @pytest.mark.parametrize( "batch_size", (1, 32), ) @pytest.mark.parametrize( "max_seq_len", (256,), # For decode-only unit test, there's no need to run with large sequence lengths ) @pytest.mark.parametrize("use_hf_rope", (True, False), ids=("hf_rope", "mllama_rope")) @pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True) def test_attention_inference( max_seq_len, batch_size, paged_attention, page_params, mesh_device, use_hf_rope, reset_seeds, use_prefetcher, ensure_gc, ): mode = Mode.DECODE dtype = ttnn.bfloat8_b pcc = 0.99 llama90b_hf_rope_pcc = 0.97 llama33_70b_mllama_rope_pcc = 0.9891 num_tensors = 2 prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None if use_prefetcher: prefetcher.init(mode) model_args = ModelArgs( mesh_device, max_batch_size=batch_size, max_seq_len=max_seq_len, cache_hf=True, prefetcher=prefetcher, use_hf_rope=use_hf_rope, ) if model_args.model_name == "Llama-3.2-90B-Instruct" and use_hf_rope: pcc = llama90b_hf_rope_pcc elif model_args.model_name == "Llama-3.3-70B-Instruct" and not use_hf_rope: pcc = llama33_70b_mllama_rope_pcc model_args.n_layers = 1 # For the unit test, just run a single layer state_dict = model_args.load_state_dict() reference_model = model_args.reference_attention(load_checkpoint=True) seq_len = 1 generation_start_pos = 0 generation_length = 10 all_tests_pass = True DefaultRopeSetup = HfRotarySetup if model_args.use_hf_rope else RotarySetup # Setup RoPE transformation matrices rope_setup = DefaultRopeSetup( mesh_device, batch_size, model_args.head_dim, model_args.max_seq_len, model_args.rope_theta, model_args.rope_scaling, model_args.use_qk_fused, prefetcher=prefetcher, ) transformation_mats = rope_setup.get_both_trans_mats() page_table_tt = None paged_attention_config = None if paged_attention: paged_attention_config = PagedAttentionConfig( block_size=page_params["page_block_size"], max_num_blocks=page_params["page_max_num_blocks"], ) # Implied shuffling of blocks permutation = torch.randperm(paged_attention_config.max_num_blocks) # Page table which maps virtual blocks to physical reverse_permutation = torch.argsort(permutation) page_table = reverse_permutation.reshape( model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size ) page_table_tt = ttnn.from_torch( page_table, device=mesh_device, dtype=ttnn.int32, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_mapper=ttnn.ShardTensor2dMesh( mesh_device, dims=(None, -2) if (model_args.is_galaxy and batch_size > 1) else (None, None), mesh_shape=model_args.cluster_shape, ), ) tt_ccl = TT_CCL(mesh_device) tt_model = Attention( mesh_device, tt_ccl, model_args, state_dict, weight_cache_path=model_args.weight_cache_path(dtype), layer_num=0, dtype=dtype, transformation_mats=transformation_mats, configuration=model_args, paged_attention_config=paged_attention_config, prefetcher=prefetcher, ) if prefetcher is not None and mode == Mode.DECODE: prefetcher.prefetch() # Prefetcher global CB size must be set to the max tensor block size amongst all 5 matmul weights # 700 is an arbitrary value that is sufficient and avoids memory clobberring prefetcher.max_tensor_block_size = 700 * 1088 cos, sin = precompute_freqs( model_args.head_dim, model_args.max_seq_len * 2, model_args.rope_theta, model_args.rope_scaling.factor if model_args.rope_scaling else None, model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None, model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3", ) freqs_cis = torch.complex(cos, sin) # Initial positions current_pos = torch.tensor([generation_start_pos for _ in range(batch_size)]) current_pos_tensor = ttnn.from_torch( current_pos, device=mesh_device, dtype=ttnn.int32, mesh_mapper=ttnn.ShardTensor2dMesh( mesh_device, dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None), mesh_shape=model_args.cluster_shape, ), ) for i in range(generation_length): # 70B attention block typically sees tensors with mean 0 and std 0.03 - 0.05 in layer 1 pt_attention_input = torch.randn( batch_size, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name) ) # Qwen2.5 0.5B sees 0.1 to 2.1 if prefetcher is not None and mode == Mode.DECODE: prefetcher.run() tt_attention_input = pt_attention_input.clone() attention_input = model_args.prepare_residual_tensor_decode( tt_attention_input, model_args.get_attn_input_mem_config(mode, prefetcher), force_replicated=False if model_args.is_galaxy else True, ) # Get cos/sin matrices for the current position of each user # When using hf style rope, those matrix does not have user dimension, # the same position is used for all of them (see #https://github.com/tenstorrent/tt-metal/issues/38223) rot_mats = rope_setup.get_rot_mats(current_pos) tt_out = tt_model( attention_input, current_pos_tensor, rot_mats=rot_mats, mode=mode, page_table=page_table_tt, ) # multi-device attention module returns replicated output tt_out = ttnn.to_torch( tt_out, mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape), ) tt_output_torch = tt_out[:, 0:1, : model_args.max_batch_size, : model_args.dim].view(-1, 1, model_args.dim) # In this test all users have the same position (if using batch > 1) freqs_cis_i = freqs_cis[current_pos[0], :].unsqueeze(0) reference_output = reference_model(pt_attention_input, current_pos[0], freqs_cis_i, mask=None) passing, pcc_message = comp_pcc(reference_output, tt_output_torch, pcc) logger.info(comp_allclose(reference_output, tt_output_torch)) logger.info(f"PCC: {pcc_message}") if passing: logger.info(f"[pos={current_pos[0]}] Attention Passed!") else: logger.warning(f"[pos={current_pos[0]}] Attention Failed!") all_tests_pass = False # Increment position current_pos = torch.tensor([generation_start_pos + i + 1 for _ in range(batch_size)]) current_pos_tensor = ttnn.from_torch( current_pos, device=mesh_device, dtype=ttnn.int32, mesh_mapper=ttnn.ShardTensor2dMesh( mesh_device, dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None), mesh_shape=model_args.cluster_shape, ), ) check_kv_cache = True if check_kv_cache: # PyTorch output -------------------------------------------------------------------- pytorch_layer_present = [ reference_model.cache_k.clone().permute(0, 2, 1, 3), # [batch_size, n_kv_heads, seq, head_dim] reference_model.cache_v.clone().permute(0, 2, 1, 3), # [batch_size, n_kv_heads, seq, head_dim] ] # TT hardware execution ------------------------------------------------------------- if paged_attention: tt_layer_present = [ ( ttnn.to_torch( cache, mesh_composer=ttnn.ConcatMesh2dToTensor( mesh_device, dims=(1, 3) if model_args.is_galaxy else (0, 1), mesh_shape=model_args.cluster_shape, ), )[reverse_permutation][:, : model_args.n_kv_heads, :, : model_args.head_dim] .reshape( model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size, model_args.n_kv_heads, paged_attention_config.block_size, model_args.head_dim, ) .transpose(1, 2) .reshape(model_args.max_batch_size, model_args.n_kv_heads, -1, model_args.head_dim)[ :batch_size, ... ] ) for cache in tt_model.layer_past ] else: tt_layer_present = [ ttnn.to_torch( cache, mesh_composer=ttnn.ConcatMesh2dToTensor( mesh_device, dims=(1, 0) if model_args.is_galaxy else (0, 1), mesh_shape=model_args.cluster_shape, ), )[:batch_size, :, :, :] for cache in tt_model.layer_past ] for label, cache_pt, cache_tt in zip(["K", "V"], pytorch_layer_present, tt_layer_present): cache_length_to_check = min(model_args.max_seq_len, generation_start_pos + i + 1) cache_pt = cache_pt[:, :, generation_start_pos:cache_length_to_check, :] cache_tt = cache_tt[:, :, generation_start_pos:cache_length_to_check, :] does_pass, output_pcc = comp_pcc(cache_pt, cache_tt, pcc) logger.info(f"{label} cache output: {output_pcc}") if does_pass: logger.info(f"{label} cache Passed!") else: logger.warning(f"{label} Cache Failed! PCC value is lower than {pcc}") all_tests_pass = False if all_tests_pass: logger.info("Attention output Passed!") else: logger.warning("Attention output Failed!") assert all_tests_pass, f"PCC value is lower than {pcc} for some of the outputs. Check Warnings!"