For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

apply_penalties_to_logits

def apply_penalties_to_logits[logit_type: DType, penalty_type: DType, //, target: StringSlice[ImmStaticOrigin]](logits: TileTensor[logit_type, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type], compressed_frequency_data: TileTensor[DType.int32, Storage=compressed_frequency_data.Storage, address_space=compressed_frequency_data.address_space, linear_idx_type=compressed_frequency_data.linear_idx_type], frequency_offsets: TileTensor[DType.uint32, Storage=frequency_offsets.Storage, address_space=frequency_offsets.address_space, linear_idx_type=frequency_offsets.linear_idx_type], frequency_penalty: TileTensor[penalty_type, Storage=frequency_penalty.Storage, address_space=frequency_penalty.address_space, linear_idx_type=frequency_penalty.linear_idx_type], presence_penalty: TileTensor[penalty_type, Storage=presence_penalty.Storage, address_space=presence_penalty.address_space, linear_idx_type=presence_penalty.linear_idx_type], repetition_penalty: TileTensor[penalty_type, Storage=repetition_penalty.Storage, address_space=repetition_penalty.address_space, linear_idx_type=repetition_penalty.linear_idx_type], ctx: DeviceContext)

Apply penalties to the logits based on the frequency of the tokens in the batch.

The frequency data is stored in a CSR format, where the frequency_offsets is the starting index of each sequence in the frequency_data array. The frequency_data array is a 2D array, where:

frequency_data[i, 0] is the token id
frequency_data[i, 1] is the frequency of the token in the sequence

Parameters:

logit_type (DType): Element type of the logits tensor (inferred).
penalty_type (DType): Element type of the per-batch penalty tensors (inferred).
target (StringSlice[ImmStaticOrigin]): Target device to dispatch the elementwise kernel to.

Args:

logits (TileTensor[logit_type, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type]): 2D logits tensor of shape [batch, vocab] updated in place with the applied penalties.
compressed_frequency_data (TileTensor[DType.int32, Storage=compressed_frequency_data.Storage, address_space=compressed_frequency_data.address_space, linear_idx_type=compressed_frequency_data.linear_idx_type]): 2D CSR frequency data where column 0 is the token id and column 1 is the token count within the sequence.
frequency_offsets (TileTensor[DType.uint32, Storage=frequency_offsets.Storage, address_space=frequency_offsets.address_space, linear_idx_type=frequency_offsets.linear_idx_type]): 1D tensor of starting indices into compressed_frequency_data for each sequence in the batch.
frequency_penalty (TileTensor[penalty_type, Storage=frequency_penalty.Storage, address_space=frequency_penalty.address_space, linear_idx_type=frequency_penalty.linear_idx_type]): Per-batch scalar multiplied by the token count and subtracted from each seen token's logit.
presence_penalty (TileTensor[penalty_type, Storage=presence_penalty.Storage, address_space=presence_penalty.address_space, linear_idx_type=presence_penalty.linear_idx_type]): Per-batch scalar subtracted from each seen token's logit.
repetition_penalty (TileTensor[penalty_type, Storage=repetition_penalty.Storage, address_space=repetition_penalty.address_space, linear_idx_type=repetition_penalty.linear_idx_type]): Per-batch scalar dividing positive logits and multiplying non-positive logits of seen tokens; must be finite and positive.
ctx (DeviceContext): Device context used to dispatch the kernel.