# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. # # SPDX-License-Identifier: Apache-2.0 """ On-device penalties module with persistent buffers, mirroring TTSampling. """ from __future__ import annotations from dataclasses import dataclass from typing import Any, List, Optional import torch import ttnn from models.common.lightweightmodule import LightweightModule @dataclass class PenaltyContext: prompt_mask: ttnn.Tensor output_mask: ttnn.Tensor output_counts: ttnn.Tensor output_counts_gathered: ttnn.Tensor presence_penalties: ttnn.Tensor frequency_penalties: ttnn.Tensor repetition_penalties: ttnn.Tensor inverse_repetition_penalties: ttnn.Tensor sub_core_grids: Any | None = None def apply_penalties(logits: ttnn.Tensor, context: Optional[PenaltyContext]) -> ttnn.Tensor: if context is None: return logits op_kwargs = {"sub_core_grids": context.sub_core_grids} if context.sub_core_grids else {} # presence presence_term = ttnn.multiply( ttnn.typecast(context.output_mask, ttnn.bfloat16, **op_kwargs), context.presence_penalties, **op_kwargs ) presence_term_bf16 = ttnn.typecast(presence_term, ttnn.bfloat16, **op_kwargs) logits = ttnn.subtract(logits, presence_term_bf16, output_tensor=logits, **op_kwargs) presence_term_bf16.deallocate() # frequency output_counts_bf16 = ttnn.typecast(context.output_counts, ttnn.bfloat16, **op_kwargs) freq_term = ttnn.multiply(output_counts_bf16, context.frequency_penalties, **op_kwargs) freq_term_bf16 = ttnn.typecast(freq_term, ttnn.bfloat16, **op_kwargs) logits = ttnn.subtract(logits, freq_term_bf16, output_tensor=logits, **op_kwargs) freq_term_bf16.deallocate() # repetition # If token appears in prompt or output, apply, otherwise use 1.0 for no-op. combined_mask_int32 = ttnn.add(context.prompt_mask, context.output_mask, **op_kwargs) combined_mask = ttnn.typecast(combined_mask_int32, ttnn.bfloat16, **op_kwargs) combined_mask_int32.deallocate() penalties = ttnn.where(combined_mask, context.repetition_penalties, 1.0, **op_kwargs) inverse_penalties = ttnn.where(combined_mask, context.inverse_repetition_penalties, 1.0, **op_kwargs) combined_mask.deallocate() # If logits are >0, divide by penalty, otherwise multiply by penalty. logits_bf16 = ttnn.typecast(logits, ttnn.bfloat16, **op_kwargs) logits_gt1 = ttnn.gt(logits_bf16, 0, **op_kwargs) scaling = ttnn.where(logits_gt1, inverse_penalties, penalties, **op_kwargs) logits_gt1.deallocate() penalties.deallocate() inverse_penalties.deallocate() logits = ttnn.multiply(logits, scaling, output_tensor=logits, **op_kwargs) scaling.deallocate() return logits class TTPenalties(LightweightModule): """ Penalty module with persistent device tensors, similar to TTSampling. """ def __init__(self, mesh_device, args): super().__init__() self.mesh_device = mesh_device self.cluster_shape = mesh_device.shape # Floor at 32 so that ROW_MAJOR [batch, vocab] buffers passed to # ttnn.tilize always have physical_volume divisible by TILE_HW # (32*32 = 1024). 32 * V is 1024-aligned for any 32-aligned V. self.max_batch_size = max(getattr(args, "max_batch_size", 32), 32) padded_vocab_size = getattr(args, "padded_vocab_size", None) self.vocab_size = padded_vocab_size if padded_vocab_size is not None else args.vocab_size self.sub_core_grids = getattr(args, "sub_core_grids", None) self._op_kwargs = {"sub_core_grids": self.sub_core_grids} if self.sub_core_grids else {} # sampling_dp > 1 when multiple mesh rows each sample independently # (e.g. GPT-OSS on [4,8] Galaxy: 4 rows × 32 users = 128 total) self._sampling_dp = getattr(args, "sampling_dp", 1) # When rows are used for data parallelism (sampling_dp > 1), vocab # must be sharded along columns; otherwise pick the larger dimension. if self._sampling_dp > 1: num_devices = mesh_device.shape[-1] else: num_devices = max(mesh_device.shape[-1], mesh_device.shape[-2]) self.num_devices = num_devices # Total batch across all rows. Host tensors use this size; after # (0, ...) sharding each row gets max_batch_size entries. self._total_batch = self.max_batch_size * self._sampling_dp # shard vocab size over larger cluster dim if mesh_device.shape[-1] == self.num_devices: shard_dims = (None, 1) shard_dims_slice = (None, 0) else: shard_dims = (1, None) shard_dims_slice = (0, None) # For row-sharded mode (sampling_dp > 1), also shard the batch dimension # across mesh rows so each row gets its own per-user penalty state. if self._sampling_dp > 1: assert ( mesh_device.shape[-1] == self.num_devices ), "Row-sharded penalties require vocab sharding along mesh columns" shard_dims = (0, 1) # batch across rows, vocab across cols shard_dims_gathered = (0, None) # batch across rows, vocab replicated shard_dims_bf16 = (0, None) # per-row penalty params per_row_batch = self.max_batch_size # NOT divided: each row gets max_batch_size else: shard_dims_gathered = (None, None) shard_dims_bf16 = None per_row_batch = self.max_batch_size self.per_row_batch_size = per_row_batch self._shard_dims_gathered = shard_dims_gathered self.prompt_mask = self._alloc_int_buffer(shard_dims=shard_dims) # Host shadow of the per-slot prompt tokens, so a partial update keeps other rows' masks. self._prompt_tokens_host = None self.output_mask = self._alloc_int_buffer(shard_dims=shard_dims) self.output_counts_gathered = self._alloc_int_buffer(shard_dims=shard_dims_gathered) self.output_counts = self._alloc_int_buffer(shard_dims=shard_dims) self._shard_dims_mask = shard_dims self.decode_src = self._alloc_int_buffer( host=torch.ones(self._total_batch, 1), shard_dims=shard_dims_gathered, layout=ttnn.ROW_MAJOR_LAYOUT ) self.zeros = self._alloc_int_buffer(shard_dims=shard_dims_gathered, layout=ttnn.ROW_MAJOR_LAYOUT) self.presence_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16) self.frequency_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16) self.repetition_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16) self.inverse_repetition_penalties = self._alloc_bf16_buffer(shard_dims=shard_dims_bf16) vocab_per_dev = self.vocab_size // self.num_devices d = torch.arange(self.num_devices, dtype=torch.int32) # [0, 0, 0, vocab_per_dev, 0, 2*vocab_per_dev, ...] start_1d = torch.empty(2 * self.num_devices, dtype=torch.int32) start_1d[0::2] = 0 start_1d[1::2] = d * vocab_per_dev # [batch, vocab_per_dev, batch, 2*vocab_per_dev, ...] end_1d = torch.empty(2 * self.num_devices, dtype=torch.int32) end_1d[0::2] = per_row_batch # per-row batch size, exclusive end_1d[1::2] = (d + 1) * vocab_per_dev # exclusive self.slice_start = ttnn.from_torch( start_1d, device=self.mesh_device, mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims_slice, mesh_shape=self.cluster_shape), ) self.slice_end = ttnn.from_torch( end_1d, device=self.mesh_device, mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims_slice, mesh_shape=self.cluster_shape), ) def _alloc_int_buffer(self, shard_dims, host=None, layout=ttnn.TILE_LAYOUT): if host is None: host = torch.zeros((self._total_batch, self.vocab_size), dtype=torch.int32) return ttnn.from_torch( host, dtype=ttnn.int32, layout=layout, device=self.mesh_device, mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape), memory_config=ttnn.DRAM_MEMORY_CONFIG, ) def _alloc_bf16_buffer(self, shard_dims=None): host = torch.zeros((self._total_batch, 1), dtype=torch.float32) if shard_dims is not None: return ttnn.from_torch( host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.mesh_device, mesh_mapper=ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape), ) return ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=self.mesh_device) def _copy_host_to_device(self, dst: ttnn.Tensor, src: torch.Tensor): if self._sampling_dp > 1: # For row-sharded buffers, create a properly sharded host tensor # so copy_host_to_device_tensor writes per-row shards correctly. mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=(0, None), mesh_shape=self.cluster_shape) src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None, mesh_mapper=mapper) else: src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None) ttnn.copy_host_to_device_tensor(src_tt, dst) def _copy_int_host_to_device(self, dst: ttnn.Tensor, src: torch.Tensor, shard_dims): mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=shard_dims, mesh_shape=self.cluster_shape) src_tt = ttnn.from_torch(src, dtype=dst.dtype, layout=ttnn.TILE_LAYOUT, device=None, mesh_mapper=mapper) ttnn.copy_host_to_device_tensor(src_tt, dst) def _token_counts_host(self, tokens_2d: torch.Tensor) -> torch.Tensor: valid = (tokens_2d >= 0) & (tokens_2d < self.vocab_size) token_ids = torch.where(valid, tokens_2d, torch.zeros_like(tokens_2d)).to(torch.int64) counts = torch.zeros((self._total_batch, self.vocab_size), dtype=torch.int32) counts.scatter_add_(1, token_ids, valid.to(torch.int32)) return counts def reset_params(self, presence: List[float], frequency: List[float], repetition: List[float]): presence_tensor = self._pad_params(presence) frequency_tensor = self._pad_params(frequency) repetition_tensor = self._pad_params(repetition) inverse_repetition_tensor = 1 / repetition_tensor self._copy_host_to_device(self.presence_penalties, presence_tensor) self._copy_host_to_device(self.frequency_penalties, frequency_tensor) self._copy_host_to_device(self.repetition_penalties, repetition_tensor) self._copy_host_to_device(self.inverse_repetition_penalties, inverse_repetition_tensor) def _pad_params(self, values: List[float]) -> torch.Tensor: tensor = torch.tensor(values, dtype=torch.float32) if tensor.numel() < self._total_batch: pad_value = tensor[-1] if tensor.numel() > 0 else torch.tensor(0.0) pad = pad_value.repeat(self._total_batch - tensor.numel()) tensor = torch.cat([tensor, pad]) elif tensor.numel() > self._total_batch: tensor = tensor[: self._total_batch] return tensor.view(self._total_batch, 1) def _pad_batch_to_max(self, tokens_2d: torch.Tensor, pad_value: int) -> torch.Tensor: """Pad/truncate first dim to _total_batch.""" if tokens_2d.dim() != 2: raise ValueError(f"Expected 2D tensor [B, S], got {tokens_2d.shape}") B, S = tokens_2d.shape if B < self._total_batch: pad = torch.full((self._total_batch - B, S), pad_value, dtype=tokens_2d.dtype) return torch.cat([tokens_2d, pad], dim=0) if B > self._total_batch: return tokens_2d[: self._total_batch] return tokens_2d def reset_prompt_tokens(self, prompt_tokens: torch.Tensor, slots: list[int] | None = None): """Rebuild the prompt mask. With ``slots``, only those rows are taken from ``prompt_tokens``; every other row keeps the prompt it was last given. The device buffer covers all rows at once, so a caller that only knows about the requests it is prefilling used to zero everyone else's mask: rows outside the call arrive as the -1 padding and hash to an empty mask. repetition_penalty is the only consumer of prompt_mask, so a live request silently stopped penalising its own prompt until something refreshed it. A later full sampling-state reset can hide this bug. A demo without that reset keeps the wiped mask for the rest of the generation. """ prompt_tokens_2d = prompt_tokens.reshape(-1, prompt_tokens.shape[-1]) prompt_tokens_2d = self._pad_batch_to_max(prompt_tokens_2d, pad_value=-1) if slots is None: self._prompt_tokens_host = prompt_tokens_2d.clone() else: shadow = getattr(self, "_prompt_tokens_host", None) width = max(prompt_tokens_2d.shape[-1], shadow.shape[-1] if shadow is not None else 0) merged = torch.full((self._total_batch, width), -1, dtype=prompt_tokens_2d.dtype) if shadow is not None: merged[:, : shadow.shape[-1]] = shadow for slot in slots: slot = int(slot) if 0 <= slot < self._total_batch: merged[slot, :] = -1 merged[slot, : prompt_tokens_2d.shape[-1]] = prompt_tokens_2d[slot] self._prompt_tokens_host = merged prompt_tokens_2d = merged # Build reset masks on host to avoid device scatter_add races on # duplicate prompt token ids (common in penalty tests/prompts). prompt_counts = self._token_counts_host(prompt_tokens_2d) prompt_mask = (prompt_counts > 0).to(torch.int32) self._copy_int_host_to_device(self.prompt_mask, prompt_mask, self._shard_dims_mask) def reset_output_tokens(self, tokens=None, slots: list[int] | None = None): if slots is not None: slots = sorted({int(slot) for slot in slots}) if any(slot < 0 or slot >= self._total_batch for slot in slots): raise ValueError(f"Output reset slots must be in [0, {self._total_batch}), got {slots}") if not slots: return # Clear only the admitted slots. A [batch, 1] device mask is # replicated across mesh columns and broadcast across vocabulary, # so continuing requests keep their accumulated device counts. keep_rows = torch.ones((self._total_batch, 1), dtype=torch.int32) keep_rows[slots] = 0 keep_rows_tt = self._alloc_int_buffer( host=keep_rows, shard_dims=self._shard_dims_gathered, ) self.output_mask = ttnn.mul( self.output_mask, keep_rows_tt, output_tensor=self.output_mask, **self._op_kwargs ) self.output_counts = ttnn.mul( self.output_counts, keep_rows_tt, output_tensor=self.output_counts, **self._op_kwargs ) self.output_counts_gathered = ttnn.mul( self.output_counts_gathered, keep_rows_tt, output_tensor=self.output_counts_gathered, **self._op_kwargs, ) keep_rows_tt.deallocate() if tokens is None: return # Restore any supplied history for the reset slots. Rows outside # ``slots`` are zero here, so adding cannot change live requests. tokens_2d = tokens.reshape(-1, tokens.shape[-1]) tokens_2d = self._pad_batch_to_max(tokens_2d, pad_value=-1) output_counts = self._token_counts_host(tokens_2d) reset_rows = torch.zeros((self._total_batch, 1), dtype=torch.int32) reset_rows[slots] = 1 output_counts *= reset_rows output_mask = (output_counts > 0).to(torch.int32) updates = ( (self.output_counts_gathered, output_counts, self._shard_dims_gathered), (self.output_counts, output_counts, self._shard_dims_mask), (self.output_mask, output_mask, self._shard_dims_mask), ) for destination, host_update, shard_dims in updates: update_tt = self._alloc_int_buffer(host=host_update, shard_dims=shard_dims) ttnn.add(destination, update_tt, output_tensor=destination, **self._op_kwargs) update_tt.deallocate() return # ALWAYS reset output buffers to zero first (this is the core accuracy fix from issue #35731) # This ensures penalty statistics are cleared between prefill and decode phases self.output_mask = ttnn.mul(self.output_mask, 0, output_tensor=self.output_mask, **self._op_kwargs) self.output_counts = ttnn.mul(self.output_counts, 0, output_tensor=self.output_counts, **self._op_kwargs) self.output_counts_gathered = ttnn.mul( self.output_counts_gathered, 0, output_tensor=self.output_counts_gathered, **self._op_kwargs ) # THEN optionally repopulate if tokens are provided if tokens is not None: tokens_2d = tokens.reshape(-1, tokens.shape[-1]) tokens_2d = self._pad_batch_to_max(tokens_2d, pad_value=-1) output_counts = self._token_counts_host(tokens_2d) output_mask = (output_counts > 0).to(torch.int32) self._copy_int_host_to_device(self.output_counts_gathered, output_counts, self._shard_dims_gathered) self._copy_int_host_to_device(self.output_counts, output_counts, self._shard_dims_mask) self._copy_int_host_to_device(self.output_mask, output_mask, self._shard_dims_mask) def update_output_tokens(self, new_tokens): # Reshape decode token to [batch, 1] for scatter_add. # Non-row-sharded: token shape is [1,1,1,batch] → shape[-1]==batch, shape[-2]==1 # Row-sharded: token shape is [1,1,batch,1] → shape[-2]==batch, shape[-1]==1 batch = self.per_row_batch_size fast_path = (new_tokens.shape[-1] == batch and new_tokens.shape[-2] == 1) or ( new_tokens.shape[-2] == batch and new_tokens.shape[-1] == 1 ) if fast_path: new_tokens = ttnn.reshape(new_tokens, [batch, 1], **self._op_kwargs) src = self.decode_src else: src = self._alloc_int_buffer( host=torch.ones(self._total_batch, new_tokens.shape[-1]), shard_dims=self._shard_dims_gathered, layout=ttnn.ROW_MAJOR_LAYOUT, ) self.token_bin_counts_and_mask( new_tokens=new_tokens, counts=self.output_counts_gathered, src=src, counts_sliced=self.output_counts, mask=self.output_mask, ) def token_bin_counts_and_mask(self, new_tokens, src, counts=None, mask=None, counts_sliced=None): counts_new = ttnn.scatter_add(self.zeros, 1, new_tokens, src, **self._op_kwargs) new_tokens.deallocate() # need to use use_low_perf because llama galaxy runs out of L1 otherwise counts_new = ttnn.tilize( counts_new, **self._op_kwargs, use_low_perf=True if self.sub_core_grids is not None else False ) if counts: counts = ttnn.add(counts, counts_new, output_tensor=counts, **self._op_kwargs) else: counts = counts_new counts_sliced = ttnn.slice( counts, self.slice_start, self.slice_end, output_tensor=counts_sliced, slice_dim=1, num_devices=self.num_devices, **self._op_kwargs, ) mask = ttnn.gt(counts_sliced, 0, output_tensor=mask, **self._op_kwargs) return counts, mask def apply(self, tt_logits: ttnn.Tensor) -> ttnn.Tensor: if tt_logits is None: return tt_logits context = PenaltyContext( prompt_mask=self.prompt_mask, output_mask=self.output_mask, output_counts=self.output_counts, output_counts_gathered=self.output_counts_gathered, presence_penalties=self.presence_penalties, frequency_penalties=self.frequency_penalties, repetition_penalties=self.repetition_penalties, inverse_repetition_penalties=self.inverse_repetition_penalties, sub_core_grids=self.sub_core_grids, ) original_shape = tt_logits.shape reshaped = ttnn.reshape(tt_logits, (-1, original_shape[-1])) apply_penalties(reshaped, context) return ttnn.reshape(reshaped, original_shape)