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

Python function

distributed_logits_postprocess

distributed_logits_postprocess()

max.nn.transformer.distributed_logits_postprocess(h, input_row_offsets, return_n_logits, lm_head, signal_buffers, return_logits, device, norm_shards=None, return_hidden_states=ReturnHiddenStates.NONE, logits_scaling=1.0, logit_softcapping=None, capture_hidden_states=None)

source

Common logits postprocessing for multi-device sharded models.

Handles last-token gathering, logits computation (VARIABLE/ALL/LAST_TOKEN), logits scaling, and hidden states return for models that use per-device sharded hidden states.

Parameters:

• h (Sequence[TensorValue]) – Per-device hidden states from the final transformer layer.
• input_row_offsets (Sequence[TensorValue]) – Per-device row offsets for ragged batching.
• return_n_logits (TensorValue) – Number of logits to return per sequence.
• lm_head (Callable[[list[TensorValue], Sequence[BufferValue]], Sequence[TensorValue]]) – Language model head (takes per-device inputs + signal buffers).
• signal_buffers (Sequence[BufferValue]) – Signal buffers for collective operations.
• return_logits (ReturnLogits) – Which logits to return.
• device (DeviceRef) – Primary device for scalar ops (such as ops.range).
• norm_shards (Sequence[Callable[[TensorValue], TensorValue]] | None) – Per-device normalization functions. When None, hidden states are passed directly to lm_head without normalization.
• return_hidden_states (ReturnHiddenStates) – Which hidden states to return.
• logits_scaling (float) – Scaling factor for logits.
• capture_hidden_states (list[list[TensorValue]] | None) – For SELECTED_LAYERS mode, the per-layer captured hidden states gathered during the layer loop. See extract_hs().
• logit_softcapping (float | None)

Returns:

Tuple of (last_logits, [logits, offsets], [hidden_states]).

Return type:

tuple[TensorValue, …]