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

get_seq_info

def get_seq_info[MaxSeqLenType: OptionallyStaticInt, ValidLengthType: OptionalPointer, PartitionType: MHAPartitionScheme, //, BM: Int, num_heads: Int, flip_prompt_idx: Bool, pair_cta: Bool = False, splitk_partitions: UInt32 = UInt32(1)](batch_size: UInt32, max_seq_len: MaxSeqLenType, valid_length: ValidLengthType, partition: PartitionType) -> SeqInfo

Computes the SeqInfo for the current CTA by querying the transient tile scheduler.

Parameters:

MaxSeqLenType (OptionallyStaticInt): Type of the maximum sequence length, which may be static or dynamic.
ValidLengthType (OptionalPointer): Optional pointer type for the per-batch valid-length tensor.
PartitionType (MHAPartitionScheme): KV-cache partitioning scheme.
BM (Int): Tile block size in the query (row) dimension, in elements.
num_heads (Int): Number of query attention heads.
flip_prompt_idx (Bool): Whether to reverse the prompt-index ordering.
pair_cta (Bool): Whether to schedule paired CTAs per tile (defaults to False).
splitk_partitions (UInt32): Number of split-K partitions to schedule across (defaults to 1).

Args:

batch_size (UInt32): Number of sequences in the batch.
max_seq_len (MaxSeqLenType): Maximum sequence length across the batch.
valid_length (ValidLengthType): Per-batch valid (non-padded) sequence lengths.
partition (PartitionType): KV-cache partition descriptor for the current launch.

Returns:

SeqInfo