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 struct

LogProbabilitiesRagged

struct LogProbabilitiesRagged

Registers the compute_log_probabilities_ragged graph op computing per-token log probabilities over ragged batches.

Implemented traits

AnyType, ImplicitlyDeletable

Methods

execute

static def execute[target: StringSlice[ImmStaticOrigin], levels: Int](lp_logits: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=lp_logits.static_spec], lp_tokens: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=lp_tokens.static_spec], logits: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logits.static_spec], tokens: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=tokens.static_spec], sampled_tokens: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=sampled_tokens.static_spec], logit_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logit_row_offsets.static_spec], token_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_row_offsets.static_spec], lp_output_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets.static_spec], lp_output_offsets_host: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets_host.static_spec], ctx: DeviceContext)

Executes the ragged log-probabilities computation across all batches.

Parameters:

• ​target (StringSlice[ImmStaticOrigin]): Compilation target string, such as "cpu" or "gpu".
• ​levels (Int): Number of heap levels controlling the top-k width; the output stores 2**levels entries per token.

Args:

• ​lp_logits (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=lp_logits.static_spec]): Output tensor of shape [num_output_tokens, 2**levels] receiving the log probabilities.
• ​lp_tokens (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=lp_tokens.static_spec]): Output tensor of shape [num_output_tokens, 2**levels] receiving the token ids paired with lp_logits.
• ​logits (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logits.static_spec]): Input logits ragged by batch, shape [total_rows, vocab_size].
• ​tokens (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=tokens.static_spec]): Previously generated tokens across all batches, ragged.
• ​sampled_tokens (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=sampled_tokens.static_spec]): Most recently sampled token for each batch.
• ​logit_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logit_row_offsets.static_spec]): Per-batch start offsets into the first axis of logits.
• ​token_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_row_offsets.static_spec]): Per-batch start offsets into tokens.
• ​lp_output_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets.static_spec]): Per-batch start offsets into the output row axis.
• ​lp_output_offsets_host (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets_host.static_spec]): Host-resident copy of lp_output_offsets whose last element gives the total number of output tokens.
• ​ctx (DeviceContext): Device context used to enqueue the computation.

Raises:

Error: If lp_logits and lp_tokens disagree on axis 0, or if either output's axis 1 does not match 2**levels.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit