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.legacy.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.
• logit_softcapping (float | None) – Cap for the softcap applied to the logits as cap * tanh(logits / cap); None skips it.
• 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().

Returns:

The named outputs; see logits_outputs().

Return type:

DecoderOutputs