# SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import math import torch import ttnn from models.common.lightweightmodule import LightweightModule from models.common.rmsnorm import RMSNorm from models.common.utility_functions import copy_to_buffer, nearest_32 from models.tt_transformers.tt.ccl import tt_all_gather, tt_all_reduce from models.tt_transformers.tt.common import Mode from models.tt_transformers.tt.model_config import OpGroup, TensorGroup, num_to_corerange class Attention(LightweightModule): def __init__( self, mesh_device, tt_ccl, args, state_dict, weight_cache_path, layer_num, dtype, transformation_mats, configuration, paged_attention_config=None, use_paged_kv_cache=False, prefetcher=None, ): super().__init__() self.args = args self.mesh_device = mesh_device self.tt_ccl = tt_ccl self.num_devices = configuration.num_devices self.prefetcher = prefetcher self.TG = self.num_devices == 32 self.hidden_size = configuration.dim self.n_heads = configuration.n_heads self.head_dim = configuration.head_dim self.max_seq_len = configuration.max_seq_len self.max_batch_size = configuration.max_batch_size self.n_kv_heads = configuration.n_kv_heads self.paged_attention_config = paged_attention_config self.min_kv_prefill_shard_seqlen = configuration.min_kv_prefill_shard_seqlen self.ccl_dtype = configuration.ccl_dtype self.MAX_QKV_MM_SEQ_LEN = configuration.MAX_QKV_MM_SEQ_LEN self.tile_size = configuration.tile_size self.rms_norm_add_unit_offset = configuration.rms_norm_add_unit_offset self.num_device_groups = self.num_devices // self.n_kv_heads self.num_devices_per_group = self.n_kv_heads if self.TG else self.num_devices self.batch_size_per_device_group = ( max(self.max_batch_size // self.num_device_groups, 1) if self.TG else self.max_batch_size ) self.n_local_heads = self.n_heads // self.num_devices_per_group self.n_local_kv_heads = self.n_kv_heads // self.num_devices_per_group self.use_qk_fused = configuration.use_qk_fused self.use_hf_rope = configuration.use_hf_rope self.arch_name = configuration.arch_name # TODO: Fix this once all-gather supports < tile_size if self.TG: weight = torch.zeros(1, 32, 8, 32) for i in range(32): col = i % 4 # This determines which group of 8 to select weight[:, i, :, col * 8 : (col + 1) * 8] = torch.eye(8) self.slice_mat = ttnn.from_torch( weight, dtype=ttnn.bfloat4_b, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, mesh_mapper=ttnn.ShardTensorToMesh(self.mesh_device, dim=1), ) user_selection_matrix = torch.eye(8, 8) user_selection_matrix = torch.nn.functional.pad(user_selection_matrix, (0, 24), "constant", 0) # (8, 32) user_selection_matrix = [user_selection_matrix] * 4 user_selection_matrix = torch.block_diag(*user_selection_matrix) # (32, 128) self.user_selection_matrix = ttnn.from_torch( user_selection_matrix, dtype=ttnn.bfloat4_b, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), ) self.dtype = dtype self.max_seq_len = configuration.max_seq_len self.grid_size = configuration.max_grid_size self.compute_kernel_config_hifi2 = configuration.compute_kernel_config_hifi2 self.compute_kernel_config_hifi2_fp16 = configuration.compute_kernel_config_hifi2_fp16 self.compute_kernel_config_hifi4 = configuration.compute_kernel_config_hifi4 self.transformation_mats = transformation_mats self.is_sliding = ( configuration.layer_types[layer_num] == "sliding_attention" if configuration.layer_types else False ) self.sliding_window = configuration.sliding_window if self.is_sliding else None self.model_config = configuration.get_model_config() self.ccl_topology = configuration.ccl_topology() self.is_multichip = configuration.is_multichip # When prefetcher is enabled, use consistent dtypes across all layers to avoid # race conditions caused by different block sizes use_prefetcher = prefetcher is not None decoders_optimizations = self.args.decoders_optimizations self.activation_dtype = decoders_optimizations.get_tensor_dtype( decoder_id=layer_num, tensor=TensorGroup.ACTIVATION, prefetcher=use_prefetcher ) self.wqkv_dtype = decoders_optimizations.get_tensor_dtype( decoder_id=layer_num, tensor=TensorGroup.WQKV, prefetcher=use_prefetcher ) self.wo_dtype = decoders_optimizations.get_tensor_dtype( decoder_id=layer_num, tensor=TensorGroup.WO, prefetcher=use_prefetcher ) self.kv_cache_dtype = decoders_optimizations.get_tensor_dtype( decoder_id=layer_num, tensor=TensorGroup.KV_CACHE, prefetcher=use_prefetcher ) self.li_qkv_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.LI_QKV_DECODE, configuration=configuration ) self.sdpa_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.SDPA_DECODE, configuration=configuration ) # Some architectures (e.g. Gemma-2, head_dim=256) need the flash-decode op's # internal default compute config instead of the framework's forced # fp32_dest_acc_en/packer_l1_acc, which corrupts the cross-chunk online-softmax # reduction for large head_dim. Config-gated flag; False for every other model. if getattr(configuration, "sdpa_decode_use_default_compute_config", False): self.sdpa_decode_compute_kernel_cfg = None self.li_o_decode_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.LI_O_DECODE, configuration=configuration ) self.sdpa_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.SDPA_PREFILL, configuration=configuration ) self.li_qkv_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.LI_QKV_PREFILL, configuration=configuration ) self.li_o_prefill_compute_kernel_cfg = decoders_optimizations.get_math_fidelity( decoder_id=layer_num, op=OpGroup.LI_O_PREFILL, configuration=configuration ) layer_name = configuration.get_state_dict_prefix(self.__class__.__name__, layer_num) if configuration.dummy_weights or (weight_cache_path is None): cache_name = lambda _: None else: cache_name = lambda name: weight_cache_path / (f"{layer_name}.{name}") # Select rotary embedding implementation for decode if self.use_hf_rope and self.use_qk_fused: raise NotImplementedError("Fused QK is not implemented for HF-style rope") if self.use_hf_rope: self.rotary_embedding_decode = self._hf_rope_decode elif self.use_qk_fused: self.rotary_embedding_decode = self._mllama_rope_fused_qk_decode else: self.rotary_embedding_decode = self._mllama_rope_decode # Select rotary embedding implementation for prefill if self.use_hf_rope: self.rotary_embedding_prefill = self._hf_rope_prefill else: self.rotary_embedding_prefill = self._mllama_rope_prefill wq_str = f"{layer_name}.wq" wk_str = f"{layer_name}.wk" wv_str = f"{layer_name}.wv" wo_str = f"{layer_name}.wo" q_norm_str = f"{layer_name}.q_norm" k_norm_str = f"{layer_name}.k_norm" # Initialize bias tensors as None self.wqkv_bias_decode = None self.wqkv_bias_prefill = None # Create combined QKV bias if present in state dict if f"{wq_str}.bias" in state_dict: bias_num_devices = self.num_devices_per_group if self.TG else configuration.num_devices qkv_bias = torch.concat( [ torch.concat( [ torch.chunk(state_dict[f"{wq_str}.bias"], bias_num_devices)[i], torch.chunk(state_dict[f"{wk_str}.bias"], bias_num_devices)[i], torch.chunk(state_dict[f"{wv_str}.bias"], bias_num_devices)[i], ], dim=-1, ) for i in range(bias_num_devices) ], dim=-1, ) # Galaxy QKV matmuls shard the input dimension over mesh columns, so the # column all-reduce sums one partial result from every column. The bias is # replicated over those columns and added to each partial before the # collective; divide it by the number of contributors so the reduced # result contains exactly one copy of the model bias. if self.TG: qkv_bias = qkv_bias / configuration.cluster_shape[1] bias_mesh_mapper = ( ttnn.ShardTensor2dMesh( self.mesh_device, dims=(-1, None), mesh_shape=configuration.cluster_shape, ) if self.TG else ttnn.ShardTensorToMesh(self.mesh_device, dim=-1) ) bias_cache_suffix = f"sharded_2d_col_reduce_{configuration.cluster_shape[1]}" if self.TG else "sharded" # Prefill can use broadcasting on the bias add so wants a 1d tensor self.wqkv_bias_prefill = ttnn.as_tensor( qkv_bias, device=self.mesh_device, mesh_mapper=bias_mesh_mapper, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG, layout=ttnn.TILE_LAYOUT, cache_file_name=cache_name(f"wqkv_bias_prefill_{bias_cache_suffix}"), ) # as_tensor returns (32, dim) which is incorrect, this reshape updates the padded size to the correct size self.wqkv_bias_prefill = ttnn.reshape( self.wqkv_bias_prefill, (1, 1, 1, self.wqkv_bias_prefill.shape[-1]), (1, 1, self.wqkv_bias_prefill.shape[-2], self.wqkv_bias_prefill.shape[-1]), ) # Broadcasting does not seem to be supported inside execute_trace so expand to the whole batch size # Create a list of bias tensors for each multiple of tile_size up to max_batch_size self.wqkv_bias_decode = [] for batch_size in range( configuration.tile_size, configuration.tile_padded_batch_rows + configuration.tile_size, configuration.tile_size, ): qkv_bias_decode = qkv_bias.unsqueeze(0).expand(batch_size, -1) bias_tensor = ttnn.as_tensor( qkv_bias_decode, device=self.mesh_device, mesh_mapper=bias_mesh_mapper, dtype=ttnn.bfloat16, memory_config=ttnn.DRAM_MEMORY_CONFIG, layout=ttnn.TILE_LAYOUT, cache_file_name=cache_name(f"wqkv_bias_decode_{bias_cache_suffix}_{batch_size}"), ) self.wqkv_bias_decode.append(bias_tensor) # when splitting the devices, we need to make sure that the number of heads is divisible by the number of devices assert self.n_heads % self.num_devices_per_group == 0 assert self.n_kv_heads % self.num_devices_per_group == 0 assert configuration.qkv_size % self.num_devices_per_group == 0 assert configuration.dim % self.num_devices_per_group == 0 # wqkv: 4096 x 3072 (2 devices): width-sharded on 12 banks, 3072 over 12 banks. wqkv_mem_config = configuration.create_dram_sharded_mem_config( configuration.dim, configuration.qkv_size // configuration.num_devices ) qkv_list = [] for i in range(self.num_devices_per_group): # Chunk weights wq_selected = torch.chunk(state_dict[f"{wq_str}.weight"], self.num_devices_per_group, dim=0)[i] wk_selected = torch.chunk(state_dict[f"{wk_str}.weight"], self.num_devices_per_group, dim=0)[i] wv_selected = torch.chunk(state_dict[f"{wv_str}.weight"], self.num_devices_per_group, dim=0)[i] # Transpose the selected chunks wq = torch.transpose(wq_selected, -2, -1) wk = torch.transpose(wk_selected, -2, -1) wv = torch.transpose(wv_selected, -2, -1) qkv = torch.cat([wq, wk, wv], dim=-1) qkv_list.append(qkv) qkv_cat = torch.cat(qkv_list, dim=-1).unsqueeze(0).unsqueeze(0) self.wqkv = ttnn.as_tensor( qkv_cat, dtype=self.wqkv_dtype, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG if self.TG else wqkv_mem_config, mesh_mapper=ttnn.ShardTensor2dMesh( self.mesh_device, dims=(3, 2) if self.TG else (2, 3), mesh_shape=configuration.cluster_shape ), cache_file_name=cache_name("wqkv_sharded_2d"), ) def norm_reshard(x, norm, mode, norm_config): """Hack until RMSNorm supports height-sharded output config""" if mode == Mode.DECODE: mem_cfg = x.memory_config() x = ttnn.to_memory_config(x, ttnn.L1_MEMORY_CONFIG, dtype=x.dtype) x = norm(x, mode, norm_config=norm_config) if mode == Mode.DECODE: x = ttnn.to_memory_config(x, mem_cfg, dtype=x.dtype) return x if f"{q_norm_str}.weight" in state_dict: fn_q_norm = RMSNorm( device=self.mesh_device, dim=self.head_dim, eps=configuration.norm_eps, state_dict=state_dict, state_dict_prefix=None, # we already prefix q_norm_str weight_cache_path=None if configuration.dummy_weights else weight_cache_path, weight_dtype=ttnn.bfloat16, weight_key=q_norm_str, add_unit_offset=self.rms_norm_add_unit_offset, is_distributed=False, tt_ccl=self.tt_ccl, ) self.q_norm = lambda x, mode, norm_config: norm_reshard(x, fn_q_norm, mode, norm_config) else: self.q_norm = lambda x, mode, norm_config: x if f"{k_norm_str}.weight" in state_dict: fn_k_norm = RMSNorm( device=self.mesh_device, dim=self.head_dim, eps=configuration.norm_eps, state_dict=state_dict, state_dict_prefix=None, # we already prefix k_norm_str weight_cache_path=None if configuration.dummy_weights else weight_cache_path, weight_dtype=ttnn.bfloat16, weight_key=k_norm_str, add_unit_offset=self.rms_norm_add_unit_offset, is_distributed=False, tt_ccl=self.tt_ccl, ) self.k_norm = lambda x, mode, norm_config: norm_reshard(x, fn_k_norm, mode, norm_config) else: self.k_norm = lambda x, mode, norm_config: x # For ring topology we can use all gather matmul for wo self.use_fused_all_gather_matmul = self.args.use_fused_all_gather_matmul pt_wo = state_dict[f"{wo_str}.weight"].transpose(-1, -2).unsqueeze(0).unsqueeze(0) wo_mem_config = configuration.create_dram_sharded_mem_config( (configuration.n_heads * configuration.head_dim) // configuration.num_devices, configuration.dim ) def get_wo_mesh_mapper(): if self.use_fused_all_gather_matmul or self.TG: return ttnn.ShardTensor2dMesh( self.mesh_device, dims=(2, 3), mesh_shape=configuration.cluster_shape, ) return ttnn.ShardTensorToMesh(self.mesh_device, dim=2) if self.prefetcher is not None: self.wo_sharded_ring = ttnn.as_tensor( pt_wo, dtype=self.wo_dtype, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, memory_config=self.args.get_sharded_wo_ring_mem_config(), mesh_mapper=get_wo_mesh_mapper(), cache_file_name=(cache_name("wo_sharded_ring")), ) def get_wo_memory_config(): if self.use_fused_all_gather_matmul or self.TG: return ttnn.DRAM_MEMORY_CONFIG else: return wo_mem_config self.wo = ttnn.as_tensor( pt_wo, dtype=self.wo_dtype, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, memory_config=get_wo_memory_config(), mesh_mapper=get_wo_mesh_mapper(), cache_file_name=( cache_name("wo_width_sharded_2d") if (self.use_fused_all_gather_matmul or self.TG) else cache_name("wo") ), ) if not use_paged_kv_cache: # vLLM provides its own kv cache self.init_kv_cache(configuration, weight_cache_path) if configuration.query_pre_attn_scalar is not None: self.scale = configuration.query_pre_attn_scalar**-0.5 else: self.scale = self.head_dim**-0.5 # Insert the tensors into the prefetcher only in decode mode, we do not use prefetcher in prefill mode if self.prefetcher is not None: def register_weights(): self.prefetcher.insert_tensor(self.wqkv) self.prefetcher.insert_tensor(self.wo_sharded_ring) self.prefetcher.register_callback(register_weights) def init_kv_cache(self, configuration, weight_cache_path): """ Generates empty KV cache and pushed to device memory """ if self.paged_attention_config: cache_k = torch.zeros( ( self.paged_attention_config.max_num_blocks, self.n_local_kv_heads, self.paged_attention_config.block_size, self.head_dim, ) ) cache_v = torch.zeros( ( self.paged_attention_config.max_num_blocks, self.n_local_kv_heads, self.paged_attention_config.block_size, self.head_dim, ) ) else: cache_k = torch.zeros( ( self.batch_size_per_device_group, self.n_local_kv_heads, self.max_seq_len, self.head_dim, ) ) cache_v = torch.zeros( ( self.batch_size_per_device_group, self.n_local_kv_heads, self.max_seq_len, self.head_dim, ) ) self.layer_past = [ ttnn.as_tensor( k_or_v, dtype=self.kv_cache_dtype, layout=self.args.get_attn_weights_layout(), device=self.mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG, mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device), cache_file_name=( f"{weight_cache_path}/kvcache_{k_or_v.shape}" if weight_cache_path and not configuration.dummy_weights else None ), ) for k_or_v in [cache_k, cache_v] ] def _update_wqkv(self, tensor: ttnn.Tensor) -> None: """In-place replace ``self.wqkv`` via ``ttnn.copy``. Caller must match the constructor's ``self.wqkv``: shape ``(1, 1, H, qkv_size_per_device)``, TILE, ``ShardTensor2dMesh(dims=(2, 3))`` (or ``(3, 2)`` on TG). """ copy_to_buffer(tensor, self.wqkv, self.wqkv_dtype) def _update_wo(self, tensor: ttnn.Tensor) -> None: """In-place replace ``self.wo`` (and ``self.wo_sharded_ring`` when the prefetcher is enabled) via ``ttnn.copy``. Both must stay in sync: the decode path reads ``wo_sharded_ring`` instead of ``wo`` when ``self.prefetcher is not None``. """ copy_to_buffer(tensor, self.wo, self.wo_dtype) wo_sharded_ring = getattr(self, "wo_sharded_ring", None) if wo_sharded_ring is not None: copy_to_buffer(tensor, wo_sharded_ring, self.wo_dtype) def _update_wqkv_bias( self, q_proj_bias: ttnn.Tensor, k_proj_bias: ttnn.Tensor, v_proj_bias: ttnn.Tensor, ) -> None: """In-place replace ``self.wqkv_bias_prefill`` and ``self.wqkv_bias_decode`` from separate HF Q/K/V biases. Not yet implemented (needs the constructor's per-device slice+concat and decode batch broadcast). """ raise NotImplementedError("wqkv_bias update is not yet implemented") def _update_qk_norm(self, which: str, weight: ttnn.Tensor) -> None: """In-place replace the optional Q-/K-RMSNorm gamma. Not yet implemented: ``self.q_norm``/``self.k_norm`` are lambdas closing over an internal ``RMSNorm`` not exposed on ``self``, so we can't delegate to ``RMSNorm.update`` without a small refactor. """ raise NotImplementedError(f"{which} update is not yet implemented") def update( self, *, q_proj: ttnn.Tensor, k_proj: ttnn.Tensor, v_proj: ttnn.Tensor, o_proj: ttnn.Tensor, q_proj_bias: ttnn.Tensor | None = None, k_proj_bias: ttnn.Tensor | None = None, v_proj_bias: ttnn.Tensor | None = None, q_norm: ttnn.Tensor | None = None, k_norm: ttnn.Tensor | None = None, hf_rope: bool = False, ) -> None: """In-place replace the on-device attention weights via ``ttnn.copy``. HF-format input contract (see ``LLAMA_WEIGHT_TRANSFER.md``): required ``self_attn.{q,k,v,o}_proj.weight`` plus optional biases/QK-norm gammas, as HF Linear shapes wrapped in two leading unit dims, bf16, TILE, DRAM-interleaved, replicated. ``hf_rope`` (default ``False``): whether Q/K rows are in HF split-half rotary order or already in this model's convention. ``False`` is right for the ttml -> tt-transformers transfer (both store Meta-permuted Q/K for Llama-3.2-1B-Instruct). When it disagrees with ``self.use_hf_rope`` an HF<->Meta row permutation is required; not implemented (raises). Conversion (single-device only): transpose each proj (mirrors the constructor's ``torch.transpose``), concat Q/K/V along the output dim, then ``_update_wqkv``/``_update_wo`` ``ttnn.copy`` into existing buffers, preserving buffer addresses (so captured traces and the prefetcher's recorded addresses stay valid). The multi-device per-device chunk-transpose-concat is not yet implemented. Raises ``ValueError`` when optional inputs don't match construction-time presence (e.g. biases passed but the module has no QKV bias). """ if hf_rope != self.use_hf_rope: raise NotImplementedError( "Attention.update: HF<->Meta Q/K row permutation is not yet " f"implemented (hf_rope={hf_rope}, self.use_hf_rope={self.use_hf_rope}). " "Permute Q/K rows on the caller side, or extend update() with the " "on-device row-permutation kernel." ) if self.num_devices_per_group != 1: raise NotImplementedError( "Attention.update for num_devices_per_group > 1 is not yet " "implemented; the multi-device path needs the constructor's " "per-device chunk-transpose-concat (see __init__'s qkv_list loop)." ) q_t = ttnn.transpose(q_proj, -2, -1) k_t = ttnn.transpose(k_proj, -2, -1) v_t = ttnn.transpose(v_proj, -2, -1) wqkv_fused = ttnn.concat([q_t, k_t, v_t], dim=-1) self._update_wqkv(wqkv_fused) wo_internal = ttnn.transpose(o_proj, -2, -1) self._update_wo(wo_internal) bias_present = (q_proj_bias is not None) or (k_proj_bias is not None) or (v_proj_bias is not None) if bias_present: if self.wqkv_bias_prefill is None: raise ValueError( "Q/K/V projection biases were supplied but Attention was " "constructed without a QKV bias." ) if q_proj_bias is None or k_proj_bias is None or v_proj_bias is None: raise ValueError("Q/K/V projection biases must all be supplied together.") self._update_wqkv_bias(q_proj_bias, k_proj_bias, v_proj_bias) else: if self.wqkv_bias_prefill is not None: raise ValueError( "Attention was constructed with a QKV bias but Q/K/V " "projection biases were not supplied." ) if q_norm is not None: self._update_qk_norm("q_norm", q_norm) if k_norm is not None: self._update_qk_norm("k_norm", k_norm) def to_qk_fused_memory_config(self, q_tensor: ttnn.Tensor, k_tensor: ttnn.Tensor): """ Convert Q and K tensors to height-sharded memory layouts suitable for fused QK ops such as rotary_embedding_llama_fused_qk and the subsequent QK matmul/attention score computation. This function: - Infers the number of Q heads and KV heads from the input tensors - Shards Q and K along the batch dimension using HEIGHT sharding - Places Q and K on disjoint core regions to avoid overlap within sub_core_grids - Uses row-major shard orientation with explicit shard shapes The resulting memory layouts are compatible with fused attention kernels that expect Q and K to be distributed across separate core ranges while preserving per-head contiguity. Args: q_tensor (ttnn.Tensor): Query tensor with shape [..., batch, num_q_heads, head_dim]. k_tensor (ttnn.Tensor): Key tensor with shape [..., batch, num_kv_heads, head_dim]. sub_core_grids (ttnn.CoreRangeSet): The available core grids to place Q and K tensors on. Returns: Tuple[ttnn.Tensor, ttnn.Tensor]: (q_tensor, k_tensor) converted to sharded memory configurations. """ n_q_heads = q_tensor.shape[2] n_kv_heads = k_tensor.shape[2] q_batch = q_tensor.shape[1] k_batch = k_tensor.shape[1] assert q_batch == k_batch row_size = 8 # We assume a row size of 8 cores k_start_core = ttnn.CoreCoord(q_batch % row_size, q_batch // row_size) q_core_grid = ttnn.CoreRangeSet({num_to_corerange(q_batch)}) k_core_grid = ttnn.CoreRangeSet({num_to_corerange(k_batch, start_core=k_start_core)}) q_mem_config = ttnn.create_sharded_memory_config( shape=(nearest_32(n_q_heads), self.head_dim), core_grid=q_core_grid, strategy=ttnn.ShardStrategy.HEIGHT, orientation=ttnn.ShardOrientation.ROW_MAJOR, use_height_and_width_as_shard_shape=True, ) k_mem_config = ttnn.create_sharded_memory_config( shape=(nearest_32(n_kv_heads), self.head_dim), core_grid=k_core_grid, strategy=ttnn.ShardStrategy.HEIGHT, orientation=ttnn.ShardOrientation.ROW_MAJOR, use_height_and_width_as_shard_shape=True, ) q_tensor = ttnn.to_memory_config(q_tensor, q_mem_config) k_tensor = ttnn.to_memory_config(k_tensor, k_mem_config) return q_tensor, k_tensor def _mllama_rope_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos): # Q Rotary Embeddings q_heads_1BQD = ttnn.experimental.rotary_embedding_llama( q_heads_pre_rot_1BQD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"], is_decode_mode=True ) # K Rotary Embeddings k_heads_1BKD = ttnn.experimental.rotary_embedding_llama( k_heads_pre_rot_1BKD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"], is_decode_mode=True ) return q_heads_1BQD, k_heads_1BKD def _mllama_rope_fused_qk_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos): q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD = self.to_qk_fused_memory_config( q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD ) q_heads_1BQD, k_heads_1BKD = ttnn.experimental.rotary_embedding_llama_fused_qk( q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats[0], rot_mats[1], self.transformation_mats["decode"] ) return q_heads_1BQD, k_heads_1BKD def _hf_rope_decode(self, q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos): if q_heads_pre_rot_1BQD.dtype != ttnn.bfloat16: q_heads_pre_rot_1BQD = ttnn.typecast(q_heads_pre_rot_1BQD, dtype=ttnn.bfloat16) if k_heads_pre_rot_1BKD.dtype != ttnn.bfloat16: k_heads_pre_rot_1BKD = ttnn.typecast(k_heads_pre_rot_1BKD, dtype=ttnn.bfloat16) q_heads_1BQD = ttnn.experimental.rotary_embedding_hf( q_heads_pre_rot_1BQD, rot_mats[0], rot_mats[1], is_decode_mode=True, ) k_heads_1BKD = ttnn.experimental.rotary_embedding_hf( k_heads_pre_rot_1BKD, rot_mats[0], rot_mats[1], is_decode_mode=True, ) return q_heads_1BQD, k_heads_1BKD def _mllama_rope_prefill(self, q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats): q_heads_1QSD = ttnn.experimental.rotary_embedding_llama( q_heads_1QSD_pre_rot, rot_mats[0], rot_mats[1], self.transformation_mats["prefill"], is_decode_mode=False, ) k_heads_1KSD = ttnn.experimental.rotary_embedding_llama( k_heads_1KSD_pre_rot, rot_mats[0], rot_mats[1], self.transformation_mats["prefill"], is_decode_mode=False, ) return q_heads_1QSD, k_heads_1KSD def _hf_rope_prefill(self, q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats): if q_heads_1QSD_pre_rot.dtype != ttnn.bfloat16: q_heads_1QSD_pre_rot = ttnn.typecast(q_heads_1QSD_pre_rot, dtype=ttnn.bfloat16) q_heads_1QSD = ttnn.experimental.rotary_embedding_hf( q_heads_1QSD_pre_rot, rot_mats[0], rot_mats[1], is_decode_mode=False, ) if k_heads_1KSD_pre_rot.dtype != ttnn.bfloat16: k_heads_1KSD_pre_rot = ttnn.typecast(k_heads_1KSD_pre_rot, dtype=ttnn.bfloat16) k_heads_1KSD = ttnn.experimental.rotary_embedding_hf( k_heads_1KSD_pre_rot, rot_mats[0], rot_mats[1], is_decode_mode=False, ) return q_heads_1QSD, k_heads_1KSD def forward_decode(self, x: ttnn.Tensor, current_pos, rot_mats=None, page_table=None, kv_cache=None) -> ttnn.Tensor: """ x: (seq_len, 1, batch, dim) current_pos: (batch_size), current token position in the sequence for each user """ ### # QKV matmuls # Use HiFi2 for DRAM-sharded matmuls as they are otherwise flop-bound. Loses 1 bit of activation precision. ### if self.TG and self.prefetcher is None: x_interleaved = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG) xqkv_fused_sharded = ttnn.linear( x_interleaved, self.wqkv, memory_config=ttnn.DRAM_MEMORY_CONFIG, program_config=None, compute_kernel_config=self.li_qkv_decode_compute_kernel_cfg, dtype=self.ccl_dtype, ) ttnn.deallocate(x_interleaved) else: xqkv_fused_sharded = ttnn.linear( x, self.wqkv, memory_config=self.args.get_attn_qkv_mm_mem_config(Mode.DECODE, self.prefetcher), program_config=self.args.get_attn_qkv_program_config(Mode.DECODE, 1, self.prefetcher), compute_kernel_config=self.li_qkv_decode_compute_kernel_cfg, dtype=self.ccl_dtype if self.TG else self.activation_dtype or ttnn.bfloat16, global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None, sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) # FIXME: File bug against dram-sharded matmuls with bias if self.wqkv_bias_decode: # select the bias tensor based on the number of tiles in the rows # WARNING: must not change the batch size between compiling and executing a trace num_tiles = int(math.ceil(xqkv_fused_sharded.shape[-2] / self.tile_size)) xqkv_fused_sharded = xqkv_fused_sharded + self.wqkv_bias_decode[num_tiles - 1] ttnn.deallocate(x) qkv_all_reduce_mem_cfg = self.args.get_attn_qkv_all_reduce_output_mem_config( Mode.DECODE, list(self.mesh_device.shape)[1], self.prefetcher ) xqkv_fused = tt_all_reduce( xqkv_fused_sharded, self.mesh_device, self.tt_ccl, cluster_axis=1, memory_config=qkv_all_reduce_mem_cfg if qkv_all_reduce_mem_cfg is not None else xqkv_fused_sharded.memory_config(), sharded=not (self.TG and self.prefetcher is None), dtype=self.ccl_dtype, topology=self.ccl_topology, subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) if self.TG: # TODO: Slice the fused_query_key_value tensor get batch=8 xqkv_fused = ttnn.matmul( self.slice_mat, xqkv_fused, dtype=ttnn.bfloat16, memory_config=self.args.get_attn_create_head_input_mem_config(Mode.DECODE), ) else: # bfloat16 is required by nlp_create_qkv_heads_decode if self.prefetcher is None: xqkv_fused = ttnn.sharded_to_interleaved(xqkv_fused_sharded, ttnn.L1_MEMORY_CONFIG, ttnn.bfloat16) ttnn.deallocate(xqkv_fused_sharded) else: xqkv_fused = xqkv_fused_sharded # Reshape such that true unpadded batch is tracked in shape fqkv_shape = xqkv_fused.shape xqkv_fused = ttnn.reshape( xqkv_fused, (1, 1, self.batch_size_per_device_group, fqkv_shape[3]), (1, 1, 32, fqkv_shape[3]) ) ### # Reshape and rotary embeddings ### ( q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, v_heads_1BKD, ) = ttnn.experimental.nlp_create_qkv_heads_decode( xqkv_fused, num_heads=self.n_local_heads, num_kv_heads=self.n_local_kv_heads, memory_config=self.args.get_attn_create_head_output_mem_config(Mode.DECODE, self.prefetcher), ) norm_config = self.args.get_norm_config("attn", Mode.DECODE, None) q_heads_pre_rot_1BQD = self.q_norm(q_heads_pre_rot_1BQD, mode=Mode.DECODE, norm_config=norm_config) k_heads_pre_rot_1BKD = self.k_norm(k_heads_pre_rot_1BKD, mode=Mode.DECODE, norm_config=norm_config) ttnn.deallocate(xqkv_fused) # Q, K Rotary Embeddings q_heads_1BQD, k_heads_1BKD = self.rotary_embedding_decode( q_heads_pre_rot_1BQD, k_heads_pre_rot_1BKD, rot_mats, current_pos ) ttnn.deallocate(q_heads_pre_rot_1BQD) ttnn.deallocate(k_heads_pre_rot_1BKD) ### # KV update ### if kv_cache: keys = kv_cache[0] values = kv_cache[1] else: keys = self.layer_past[0] values = self.layer_past[1] # k_heads, [seqlen, n_kv_heads, bsz, head_dim] # v_heads [seqlen, n_kv_heads, bsz, head_dim] # keys, [max_batch_size, n_kv_heads // configuration.num_devices, max_seq_len, head_dim] if self.use_qk_fused: ttnn.experimental.paged_fused_update_cache( keys, k_heads_1BKD, values, v_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table ) else: ttnn.experimental.paged_update_cache( keys, k_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table ) ttnn.experimental.paged_update_cache( values, v_heads_1BKD, update_idxs_tensor=current_pos, page_table=page_table ) ttnn.deallocate(k_heads_1BKD) ttnn.deallocate(v_heads_1BKD) # NOTE: Varying the batch size will result in slightly different outputs. # For example, a prompt w/ 1 user vs, the same prompt repeated N times for N users, will produce different outputs # This is because the SDPA op in decode mode has different number of reductions depending on batch size # Which leads to slightly different outputs from attention (due to accumulated errors) sdpa_decode_prog_cfg = self.args.get_attn_sdpa_decode_program_config(self.prefetcher) if page_table is not None: attn_output_1G4D = ttnn.transformer.paged_scaled_dot_product_attention_decode( q_heads_1BQD, keys, values, page_table_tensor=page_table, cur_pos_tensor=current_pos, scale=self.scale, sliding_window_size=self.sliding_window, program_config=sdpa_decode_prog_cfg, compute_kernel_config=self.sdpa_decode_compute_kernel_cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) else: attn_output_1G4D = ttnn.transformer.scaled_dot_product_attention_decode( q_heads_1BQD, keys, values, cur_pos_tensor=current_pos, scale=self.scale, sliding_window_size=self.sliding_window, program_config=sdpa_decode_prog_cfg, compute_kernel_config=self.sdpa_decode_compute_kernel_cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG, # FIXME: why not L1 height sharded e.g. SCORES_BATCHED_MM_OUTPUT_MEMCFG? ) ttnn.deallocate(q_heads_1BQD) attn_output_11BH = ttnn.to_memory_config( attn_output_1G4D, memory_config=self.args.get_attn_sdpa_output_mem_config( Mode.DECODE, self.batch_size_per_device_group, self.prefetcher ), ) attn_output_cat = ttnn.experimental.nlp_concat_heads_decode( attn_output_11BH, num_heads=self.n_local_heads, sub_core_grids=self.prefetcher.all_worker_cores_range_set if self.prefetcher is not None else None, ) ttnn.deallocate(attn_output_11BH) ttnn.deallocate(attn_output_1G4D) if self.use_fused_all_gather_matmul or self.prefetcher is not None: attn_output_cat = ttnn.to_memory_config( attn_output_cat, self.args.get_attn_concat_heads_output_mem_config(Mode.DECODE, self.prefetcher), ) # Fused AGMM only valid for ring topology if self.ccl_topology == ttnn.Topology.Ring and self.prefetcher is None: _, dense_out_sharded = ttnn.experimental.all_gather_matmul_async( attn_output_cat, self.wo, persistent_output_buffer=None, dim=3, multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), all_gather_core_grid_offset=(0, 4), barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), num_links=self.model_config["ATTN_AGMM_CONFIG"]["num_links"], memory_config_ag=self.args.get_attn_all_gather_output_mem_config(Mode.DECODE, None), memory_config_mm=self.args.get_attn_dense_output_mem_config(Mode.DECODE, None), program_config=self.args.get_attn_all_gather_matmul_program_config(Mode.DECODE, None), compute_kernel_config=self.compute_kernel_config_hifi2, chunks_per_sync=self.model_config["ATTN_AGMM_CONFIG"]["chunks_per_sync"], num_workers_per_link=self.model_config["ATTN_AGMM_CONFIG"]["num_workers_per_link"], num_buffers_per_channel=2, subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) else: all_gather_output = ttnn.experimental.all_gather_async( attn_output_cat, persistent_output_buffer=None, dim=3, multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), num_links=1, topology=self.ccl_topology, memory_config=self.args.get_attn_all_gather_output_mem_config(Mode.DECODE, self.prefetcher), barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), chunks_per_sync=10, num_workers_per_link=2, num_buffers_per_channel=2, subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) dense_out_sharded = ttnn.linear( all_gather_output, self.wo_sharded_ring if self.prefetcher is not None else self.wo, memory_config=self.args.get_attn_dense_output_mem_config(Mode.DECODE, self.prefetcher), program_config=self.args.get_attn_all_gather_matmul_program_config(Mode.DECODE, self.prefetcher), compute_kernel_config=self.li_o_decode_compute_kernel_cfg, global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None, sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) ttnn.deallocate(all_gather_output) ttnn.deallocate(attn_output_cat) dense_out_sharded = ttnn.to_memory_config( dense_out_sharded, self.args.get_attn_dense_output_mem_config(Mode.DECODE, self.prefetcher), ) return dense_out_sharded else: attn_output = tt_all_gather( attn_output_cat, self.mesh_device, self.tt_ccl, dim=2, cluster_axis=1, memory_config=self.args.get_attn_gather_users_mem_config( Mode.DECODE, list(self.mesh_device.shape)[1], self.prefetcher ), sharded=True, subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, # dtype=self.ccl_dtype, # Running bf16 until we have SDPA output bfp8 df; otherwise we have two sharded to interleaved/interleaved to sharded conversions ) if self.TG: attn_output = ttnn.to_memory_config(attn_output, ttnn.L1_MEMORY_CONFIG) # user_selection_matrix = [1, 1, 32, 128] # user_selection_matrix @ activation -> [1, 1, 32, 128] * [1, 1, 128, 2048] -> [1, 1, 32, 2048] attn_output = ttnn.matmul( self.user_selection_matrix, attn_output, core_grid=ttnn.CoreGrid(y=4, x=8), dtype=ttnn.bfloat16, memory_config=ttnn.L1_WIDTH_SHARDED_MEMORY_CONFIG, ) # TODO: Fix this once self.TG supports dram-sharded matmuls dense_out_sharded = ttnn.linear( attn_output, self.wo, core_grid=ttnn.CoreGrid(y=4, x=8) if self.TG else None, program_config=self.args.get_attn_wo_program_config(Mode.DECODE, 1, self.prefetcher), memory_config=self.args.get_attn_wo_output_mem_config(Mode.DECODE, self.prefetcher), dtype=ttnn.bfloat8_b if self.TG else None, compute_kernel_config=self.li_o_decode_compute_kernel_cfg, global_cb=self.prefetcher.global_cb if self.prefetcher is not None else None, sub_device_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) ttnn.deallocate(attn_output_cat) # All reduce dense_out_reduced = tt_all_reduce( dense_out_sharded, self.mesh_device, self.tt_ccl, cluster_axis=0, dim=0 if (self.TG and self.hidden_size < 8192) else 3, topology=self.ccl_topology, memory_config=self.args.get_attn_all_reduce_output_mem_config( Mode.DECODE, self.hidden_size, list(self.mesh_device.shape)[0], self.prefetcher ), sharded=True, dtype=self.ccl_dtype, use_composite=True if self.hidden_size == 8192 else False, subdevice_id=self.prefetcher.worker_sub_device_id if self.prefetcher is not None else None, ) if not self.TG: dense_out_reduced = ttnn.to_memory_config( dense_out_reduced, self.args.get_attn_dense_output_mem_config(Mode.DECODE, None) ) return dense_out_reduced def forward_prefill( self, x_11SH, rot_mats, user_id: int = 0, page_table=None, chunk_page_table=None, chunk_start_idx=None, kv_cache=None, ): # For batched prefill, x_11SH has shape [B, 1, S, H] where B is batch_size # concat before QKV matmul, then reshape back to batch after batch_size = x_11SH.shape[0] if batch_size > 1: # Concatenate batch dimension into sequence for matmul compatibility x_11SH = ttnn.reshape(x_11SH, [1, 1, x_11SH.shape[-2] * x_11SH.shape[-3] * x_11SH.shape[-4], -1]) seq_len = x_11SH.shape[-2] original_seq_len = seq_len # Track original for later unpadding assert seq_len % 128 == 0 and seq_len > 0, "Seqlen must be divisible by 128" ### # QKV matmuls ### # reshaping long sequence to matmul fit on device # Pad seq_len to nearest multiple of MAX_QKV_MM_SEQ_LEN if needed if seq_len > self.MAX_QKV_MM_SEQ_LEN and seq_len % self.MAX_QKV_MM_SEQ_LEN != 0: padded_seq_len = ( (seq_len + self.MAX_QKV_MM_SEQ_LEN - 1) // self.MAX_QKV_MM_SEQ_LEN ) * self.MAX_QKV_MM_SEQ_LEN pad_len = padded_seq_len - seq_len x_11SH = ttnn.pad(x_11SH, padding=[(0, 0), (0, 0), (0, pad_len), (0, 0)], value=0.0) seq_len = padded_seq_len if seq_len > self.MAX_QKV_MM_SEQ_LEN: x_11SH = ttnn.reshape(x_11SH, [1, seq_len // self.MAX_QKV_MM_SEQ_LEN, self.MAX_QKV_MM_SEQ_LEN, -1]) if self.args.use_minimal_qkv_prefill_matmul(seq_len): xqkv_fused = ttnn.experimental.minimal_matmul( x_11SH, self.wqkv, compute_kernel_config=self.li_qkv_prefill_compute_kernel_cfg, config=self.args.get_attn_qkv_program_config(Mode.PREFILL, seq_len, None), ) else: xqkv_fused = ttnn.linear( x_11SH, self.wqkv, dtype=self.ccl_dtype if self.TG else self.activation_dtype or ttnn.bfloat16, memory_config=self.args.get_attn_qkv_mm_mem_config(Mode.PREFILL, None), compute_kernel_config=self.li_qkv_prefill_compute_kernel_cfg, program_config=self.args.get_attn_qkv_program_config(Mode.PREFILL, seq_len, None), ) # FIXME: surely ttnn.linear bias should work? if self.wqkv_bias_prefill is not None: xqkv_fused = xqkv_fused + self.wqkv_bias_prefill xqkv_fused = tt_all_reduce( xqkv_fused, self.mesh_device, self.tt_ccl, cluster_axis=1, memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=self.ccl_dtype, ) if seq_len > self.MAX_QKV_MM_SEQ_LEN: xqkv_fused = ttnn.reshape(xqkv_fused, [1, 1, seq_len, -1]) # Slice back to original seq_len if we padded earlier if original_seq_len != seq_len: xqkv_fused = xqkv_fused[:, :, :original_seq_len, :] seq_len = original_seq_len if batch_size > 1: xqkv_fused = ttnn.reshape(xqkv_fused, [batch_size, 1, seq_len // batch_size, -1]) ttnn.deallocate(x_11SH) # split qkv into heads ( q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, v_heads_1VSD, ) = ttnn.experimental.nlp_create_qkv_heads( xqkv_fused, num_heads=self.n_local_heads, num_kv_heads=self.n_local_kv_heads, transpose_k_heads=False, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) norm_config = self.args.get_norm_config("attn", Mode.PREFILL, None) q_heads_1QSD_pre_rot = self.q_norm(q_heads_1QSD_pre_rot, mode=Mode.PREFILL, norm_config=norm_config) k_heads_1KSD_pre_rot = self.k_norm(k_heads_1KSD_pre_rot, mode=Mode.PREFILL, norm_config=norm_config) ttnn.deallocate(xqkv_fused) ### # Rotary embeddings ### # Apply rotary embeddings using the selected implementation q_heads_1QSD, k_heads_1KSD = self.rotary_embedding_prefill(q_heads_1QSD_pre_rot, k_heads_1KSD_pre_rot, rot_mats) ttnn.deallocate(q_heads_1QSD_pre_rot) ttnn.deallocate(k_heads_1KSD_pre_rot) # Fill KV-Cache if kv_cache: keys_BKSD, values_BKSD = kv_cache[0], kv_cache[1] else: keys_BKSD, values_BKSD = self.layer_past[0], self.layer_past[1] k_heads_1KSD_8b = ttnn.typecast(k_heads_1KSD, dtype=keys_BKSD.dtype) ttnn.deallocate(k_heads_1KSD) # sharding k_fill to deal with update_cache memory limitation if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None: k_fill = ttnn.interleaved_to_sharded(k_heads_1KSD_8b, self.args.get_attn_kv_prefill_mem_config(seq_len)) else: k_fill = k_heads_1KSD_8b v_heads_1VSD_8b = ttnn.typecast(v_heads_1VSD, dtype=values_BKSD.dtype) ttnn.deallocate(v_heads_1VSD) # sharding v_fill to deal with update_cache memory limitation if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None: v_fill = ttnn.interleaved_to_sharded(v_heads_1VSD_8b, self.args.get_attn_kv_prefill_mem_config(seq_len)) else: v_fill = v_heads_1VSD_8b if self.TG: k_fill = self.prefill_prepare_tensor_for_kv_cache(k_fill, user_id) v_fill = self.prefill_prepare_tensor_for_kv_cache(v_fill, user_id) if page_table is not None: # In the case that the tokens have been padded along the seq len dimension, we need to fill the cache with the unpadded k/v values. # Assume that the page table does not have padding, so we can use it to get the unpadded page len. block_size = keys_BKSD.shape[2] # If chunked prefill, use chunk_page_table if given, otherwise use page_table. fill_page_table = chunk_page_table if chunk_page_table is not None else page_table if batch_size > 1: # For batched prefill, fill every slot of the padded batch separately. # k_fill/v_fill have shape [padded_batch, n_kv_heads, seq_len_per_user, head_dim] # The paged_fill_cache kernel reads batch_idx_ptr[0] for all positions, # so we must call it once per slot with that slot's K/V slice. # # IMPORTANT: this loop is captured into the prefill trace, whose identity is the # padded batch width, not the set of active slots. Looping over the active slots # (user_id) would freeze that slot set into the trace, and a later replay for a # different slot set would never write K/V for the new slots. Empty slots have # page_table rows of -1, which the kernel treats as "skip" without any write, so # filling all padded_batch rows is both safe and slot-agnostic. seq_len_per_user = k_fill.shape[2] page_len = fill_page_table.shape[1] * block_size for slot_idx in range(batch_size): # Extract this slot's K/V slice: [1, n_kv_heads, seq_len_per_user, head_dim] k_user = k_fill[slot_idx : slot_idx + 1, :, :, :] v_user = v_fill[slot_idx : slot_idx + 1, :, :, :] # Slice to page length if needed (same as single-user path) k_user_sliced = k_user[:, :, :page_len, :] if page_len < seq_len_per_user else k_user v_user_sliced = v_user[:, :, :page_len, :] if page_len < seq_len_per_user else v_user # Fill cache for this specific slot with scalar batch_idx ttnn.experimental.paged_fill_cache(keys_BKSD, k_user_sliced, fill_page_table, batch_idx=slot_idx) ttnn.experimental.paged_fill_cache(values_BKSD, v_user_sliced, fill_page_table, batch_idx=slot_idx) elif page_table is not None: # Single user path with page_table page_len = fill_page_table.shape[1] * block_size k_fill_sliced = k_fill[:, :, :page_len, :] if page_len < k_fill.shape[2] else k_fill v_fill_sliced = v_fill[:, :, :page_len, :] if page_len < v_fill.shape[2] else v_fill ttnn.experimental.paged_fill_cache(keys_BKSD, k_fill_sliced, fill_page_table, batch_idx=user_id) ttnn.experimental.paged_fill_cache(values_BKSD, v_fill_sliced, fill_page_table, batch_idx=user_id) else: # Single user path without page_table ttnn.fill_cache( keys_BKSD, k_fill, user_id % self.batch_size_per_device_group, ) ttnn.fill_cache( values_BKSD, v_fill, user_id % self.batch_size_per_device_group, ) if seq_len >= self.min_kv_prefill_shard_seqlen and not self.TG and page_table is None: ttnn.deallocate(k_fill) ttnn.deallocate(v_fill) # SDPA q_heads_1QSD_8b = ttnn.typecast(q_heads_1QSD, dtype=self.activation_dtype or ttnn.bfloat8_b) ttnn.deallocate(q_heads_1QSD) if chunk_start_idx is not None: if self.sliding_window is not None: raise NotImplementedError("Sliding window not supported for chunked prefill SDPA") if isinstance(chunk_start_idx, ttnn.Tensor): attn_output_84SD = ttnn.transformer.chunked_scaled_dot_product_attention( input_tensor_q=q_heads_1QSD_8b, input_tensor_k=keys_BKSD, input_tensor_v=values_BKSD, page_table_tensor=page_table, chunk_start_idx=None, chunk_start_idx_tensor=chunk_start_idx, compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg, program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, seq_len, 0, None), ) else: attn_output_84SD = ttnn.transformer.chunked_scaled_dot_product_attention( input_tensor_q=q_heads_1QSD_8b, input_tensor_k=keys_BKSD, input_tensor_v=values_BKSD, page_table_tensor=page_table, chunk_start_idx=chunk_start_idx, compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg, program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, seq_len, chunk_start_idx, None), ) else: # For batched prefill, the actual per-user seq_len is seq_len // batch_size # since the tensors have shape [batch_size, n_heads, seq_len_per_user, head_dim] sdpa_seq_len = seq_len // batch_size if batch_size > 1 else seq_len attn_output_84SD = ttnn.transformer.scaled_dot_product_attention( q_heads_1QSD_8b, k_heads_1KSD_8b, v_heads_1VSD_8b, is_causal=True, sliding_window_size=self.sliding_window, scale=self.scale, compute_kernel_config=self.sdpa_prefill_compute_kernel_cfg, program_config=self.args.get_attn_sdpa_program_config(Mode.PREFILL, sdpa_seq_len, None, None), ) # deallocate keys and values ttnn.deallocate(q_heads_1QSD_8b) ttnn.deallocate(k_heads_1KSD_8b) ttnn.deallocate(v_heads_1VSD_8b) # For single-user prefill, reshape to expected format for nlp_concat_heads # For batched prefill (batch_size > 1), skip this reshape - nlp_concat_heads handles [B, H, S, D] # IMPORTANT: Reshaping [B, H, S, D] to [1, H, B*S, D] BEFORE concat_heads would scramble data # because batch and sequence dimensions are separated by heads. Must reshape AFTER concat_heads. if batch_size == 1: attn_output_1QSD = ttnn.reshape(attn_output_84SD, [1, self.n_local_heads, -1, self.head_dim]) else: attn_output_1QSD = attn_output_84SD ### # Output matmul ### attn_output_11SH = ttnn.experimental.nlp_concat_heads( attn_output_1QSD, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) ttnn.deallocate(attn_output_1QSD) # For batched prefill, reshape to concatenate batch dimension into sequence # This MUST happen AFTER nlp_concat_heads to preserve correct data layout # nlp_concat_heads outputs [B, 1, S_per_user, H*D], reshape to [1, 1, B*S, H*D] if batch_size > 1: attn_output_11SH = ttnn.reshape(attn_output_11SH, [1, 1, seq_len, -1]) # reshaping long sequence to matmul fit on device if seq_len > 1024: attn_output_11SH = ttnn.reshape(attn_output_11SH, [1, seq_len // 1024, 1024, -1]) # Non fused All Gather Matmul if self.use_fused_all_gather_matmul: # is true for Ring topology attn_output_11SH = ttnn.experimental.all_gather_async( attn_output_11SH, persistent_output_buffer=None, dim=3, multi_device_global_semaphore=self.tt_ccl.get_and_cycle_ag_semaphore_handles(), num_links=1, topology=self.ccl_topology, memory_config=ttnn.DRAM_MEMORY_CONFIG, barrier_semaphore=self.tt_ccl.get_and_cycle_barrier_semaphore_handle(), chunks_per_sync=10, num_workers_per_link=2, num_buffers_per_channel=2, ) output_11SH = ttnn.linear( attn_output_11SH, self.wo, compute_kernel_config=self.li_o_prefill_compute_kernel_cfg, dtype=self.activation_dtype or ttnn.bfloat8_b, memory_config=ttnn.DRAM_MEMORY_CONFIG, program_config=self.args.get_attn_wo_program_config(Mode.PREFILL, seq_len, None), ) if seq_len > 1024: output_11SH = ttnn.reshape(output_11SH, [1, 1, seq_len, -1]) ttnn.deallocate(attn_output_11SH) # Reduce-scatter if not self.use_fused_all_gather_matmul: output_11SH = tt_all_reduce( output_11SH, self.mesh_device, self.tt_ccl, cluster_axis=0, dim=0 if self.TG else 3, topology=self.ccl_topology, memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=self.ccl_dtype, ) return output_11SH def forward( self, x, current_pos, rot_mats=None, user_id=0, mode=Mode.DECODE, page_table=None, chunk_page_table=None, chunk_start_idx=None, kv_cache=None, ): if mode == Mode.PREFILL: return self.forward_prefill( x, rot_mats, user_id, page_table=page_table, chunk_page_table=chunk_page_table, chunk_start_idx=chunk_start_idx, kv_cache=kv_cache, ) else: return self.forward_decode(x, current_pos, rot_mats, page_table=page_table, kv_cache=kv_cache) def prefill_prepare_tensor_for_kv_cache(self, key_or_value_layer, user_id): tensor_copy = ttnn.clone(key_or_value_layer) # key_or_value_layer.deallocate(True) # Get all tensors from multi-device tensor tensors = ttnn.get_device_tensors(tensor_copy) # Get only tensors from specific column chips # Get every 4th tensor starting from user_id // 8 single_column_tensors = tensors[user_id // self.batch_size_per_device_group :: 4] # Create multi-device tensor multi_device_tensor = ttnn.combine_device_tensors(tensors=single_column_tensors) return multi_device_tensor