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 class

MHAAttnKey

MHAAttnKey

class max.nn.kv_cache.MHAAttnKey(batch_size, max_prompt_length, num_partitions)

source

Bases: AttnKey

Decode dispatch metadata for multi-head attention (MHA).

Parameters:

• batch_size (int)
• max_prompt_length (int)
• num_partitions (int)

pack_into_buffer()

pack_into_buffer(device, max_cache_valid_length)

source

Returns the CPU dispatch buffer MHA decode kernels read, holding batch size, prompt width, partition count, and cache length.

Parameters:

• device (Device)
• max_cache_valid_length (int)

Return type:

Buffer