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

kda_chunk_supports_head_dims

kda_chunk_supports_head_dims()

max.nn.state_space.kda_chunk_supports_head_dims(key_head_dim, value_head_dim)

source

Reports whether kda_chunk is compiled for this head-dim pair.

Callers that can fall back to the sequential recurrence should ask this before routing to kda_chunk(), which raises on an unsupported pair. On an AMD accelerator only gfx950 runs kda_chunk, and only at (128, 128) with as many key heads as value heads (Hk == Hv), which this helper cannot see and the caller must ensure.

Parameters:

• key_head_dim (int) – Per-head key width.
• value_head_dim (int) – Per-head value width.

Returns:

True when the pair is one the op is compiled for.

Return type:

bool