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

ragged_kv_symbolic_input_tree

ragged_kv_symbolic_input_tree()

max.pipelines.lib.interfaces.ragged_kv_symbolic_input_tree(*, kv_params, device_refs, include_signal_buffers, batch_order=('tokens', 'row_offsets', 'return_n_logits'), hidden_states=None)

source

Returns the named graph input types of a standard ragged KV text model.

The tree’s leaves are the inputs in execution order (batch_order fields, signal buffers, KV leaves), and its paths name each one, so a consumer that needs to know what a graph slot holds reads the tree rather than re-deriving the order.

Parameters:

• kv_params (KVCacheParamInterface) – The KV cache parameters that supply the kv subtree.
• device_refs (list[DeviceRef]) – The devices the model runs on. The first holds the batch tensors.
• include_signal_buffers (bool) – Whether to add the allreduce signals.
• batch_order (Sequence[str]) – The per-batch inputs to include, in graph order. See ragged_batch_inputs().
• hidden_states (tuple[DType, int] | None) – The (dtype, hidden_size) of the EAGLE draft’s hidden-state input, appended to batch_order when given.

Returns:

The tree of input types, keyed batch, signals and kv.

Return type:

OrderedDict[str, object]