Python module

kernels

Helper functions for wrapping custom kv cache/attention related ops.

AttentionMaskVariant

class max.nn.kernels.AttentionMaskVariant(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

CAUSAL_MASK

CAUSAL_MASK = 'causal'

CHUNKED_CAUSAL_MASK

CHUNKED_CAUSAL_MASK = 'chunked_causal'

NULL_MASK

NULL_MASK = 'null'

SLIDING_WINDOW_CAUSAL_MASK

SLIDING_WINDOW_CAUSAL_MASK = 'sliding_window_causal'

TENSOR_MASK

TENSOR_MASK = 'tensor_mask'

MHAMaskConfig

class max.nn.kernels.MHAMaskConfig(attention_mask_variant: 'AttentionMaskVariant', positional_encoding_variant: 'PositionalEncodingVariant')

Parameters:

• attention_mask_variant (AttentionMaskVariant )
• positional_encoding_variant (PositionalEncodingVariant )

attention_mask_variant

attention_mask_variant*: AttentionMaskVariant*

positional_encoding_variant

positional_encoding_variant*: PositionalEncodingVariant*

MHAMaskVariant

class max.nn.kernels.MHAMaskVariant(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

CAUSAL_ALIBI_MASK

CAUSAL_ALIBI_MASK = '1'

CAUSAL_MASK

CAUSAL_MASK = '0'

CHUNKED_CAUSAL_MASK

CHUNKED_CAUSAL_MASK = '3'

NULL_MASK

NULL_MASK = '2'

SLIDING_WINDOW_CAUSAL_MASK

SLIDING_WINDOW_CAUSAL_MASK = '4'

PositionalEncodingVariant

class max.nn.kernels.PositionalEncodingVariant(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

ALIBI_POS

ALIBI_POS = 'alibi_pos'

NO_POS

NO_POS = 'no_pos'

causal_flash_attention_gpu()

max.nn.kernels.causal_flash_attention_gpu(q, k, v, scale)

Computes causal flash attention using GPU-optimized kernel. :param q: Query tensor of shape [batch, seq_len, num_heads, head_dim] :param k: Key tensor of shape [batch, seq_len, num_heads, head_dim] :param v: Value tensor of shape [batch, seq_len, num_heads, head_dim] :param scale: Scaling factor for attention scores

Parameters:

• q (TensorValue )
• k (TensorValue )
• v (TensorValue )
• scale (float )

Return type:

TensorValue

cross_attention_ragged()

max.nn.kernels.cross_attention_ragged(kv_params, input, input_row_offsets, kv_collection, layer_idx, mask_variant, kv_input_row_offsets, q_max_seq_len, scale, local_window_size=-1)

Computes cross attention provided the !mo.opaque KV Cache.

Notably, this materializes the attention mask (dependent on MHAMaskVariant) within the kernel. input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

attention, kv_input_row_offsets represents the KV sequence length.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• layer_idx (TensorValue )
• mask_variant (MHAMaskVariant )
• kv_input_row_offsets (TensorValue )
• q_max_seq_len (TensorValue )
• scale (float )
• local_window_size (int )

Return type:

TensorValue

dynamic_scaled_matmul()

max.nn.kernels.dynamic_scaled_matmul(a, b, a_scales, b_scales, out_type=bfloat16)

Perform a matmul of two tensors with scaling factors. Currently only supports channel-wise scaling for weights and per-token scaling for inputs.

Parameters:

• a (TensorValue ) – The first tensor to multiply.
• b (TensorValue ) – The second tensor to multiply, must be transposed.
• a_scales (TensorValue ) – The scaling factors for the first tensor.
• b_scales (TensorValue ) – The scaling factors for the second tensor.
• out_type (DType )

Returns:

The result of the matmul operation.

Return type:

TensorValue

flare_mla_decode_ragged()

max.nn.kernels.flare_mla_decode_ragged(kv_params, input, input_row_offsets, kv_collection, layer_idx, mask_variant, scale, qk_rope_dim=64)

Computes flash (self) attention provided the !mo.opaque KV Cache.

Notably, this materializes the attention mask (dependent on MHAMaskVariant) within the kernel. input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Note that this is self attention and the KV sequence length is assumed to be equal to the Q sequence length. For KV sequence length != Q sequence length, use cross_attention_ragged.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )
• mask_variant (MHAMaskVariant )
• scale (float )
• qk_rope_dim (int )

Return type:

TensorValue

flare_mla_decompress_k_cache()

max.nn.kernels.flare_mla_decompress_k_cache(kv_params, buffer_row_offsets_1d, cache_offsets_1d, buffer_length, weight, kv_collection, layer_idx, buffer_size)

This kernel decompresses the key cache by up-projecting latent representations into the KV space using a weight matrix.

The process involves:

Copying buffer_length latent vectors from the key cache into a contiguous buffer (k_latent) 2. Computing k = k_latent @ weight.T to obtain the decompressed keys

Returns:

A tensor of shape [buffer_size, weight.shape[0]] containing the decompressed keys. Note that only the first buffer_length tokens are valid.

Parameters:

• kv_params (KVCacheParams )
• buffer_row_offsets_1d (TensorValue )
• cache_offsets_1d (TensorValue )
• buffer_length (TensorValue )
• weight (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )
• buffer_size (int )

Return type:

TensorValue

flare_mla_prefill_plan()

max.nn.kernels.flare_mla_prefill_plan(kv_params, input_row_offsets, kv_collection, layer_idx, buffer_size, max_chunks=16)

This kernel plans how to process a batch of sequences with varying lengths using a fixed-size buffer.

Each sequence in the batch has some existing cached tokens and new input tokens. The kernel divides the total tokens into chunks of buffer_size.

For each chunk (iteration), it calculates:

Buffer offsets for each sequence in each chunk 2. Cache offsets for each sequence in each chunk 3. Total buffer lengths for each processing iteration

Parameters:

• kv_params (KVCacheParams )
• input_row_offsets (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )
• buffer_size (int )
• max_chunks (int )

Return type:

tuple[TensorValue, TensorValue, TensorValue]

flare_mla_prefill_ragged()

max.nn.kernels.flare_mla_prefill_ragged(kv_params, input, k, v, input_row_offsets, buffer_row_offsets, cache_offsets, kv_collection, layer_idx, mask_variant, scale, qk_rope_dim=64, prev_output=None, prev_softmax_info=None)

Performs MLA prefill. In the MLA prefill, we need to decompress the KV tensors, as we store the latent representations in the KV cache. We will decompress the KV tensors into a fixed size buffer to avoid out-of-memory errors. In case the total cache length is greater than the buffer size, we will process the attention calculation in chunks.

This MLA prefill kernel will return the output tensor for this iteration and the softmax info tensor for this iteration. Such tensors will be used by the next iteration of the MLA prefill kernel to continue the attention calculation.

Parameters:

• kv_params (KVCacheParams ) – KVCacheParams
• input (TensorValue ) – Input tensor
• k (TensorValue ) – Key tensor
• v (TensorValue ) – Value tensor
• input_row_offsets (TensorValue ) – Indicates where each batch starts and ends in input
• buffer_row_offsets (TensorValue ) – Indicates where each batch starts and ends in the buffer
• cache_offsets (TensorValue ) – Indicates where each batch starts and ends in the KV cache
• kv_collection (PagedKVCacheCollection ) – KV collection
• layer_idx (TensorValue ) – Layer index tensor
• mask_variant (MHAMaskVariant ) – Mask variant
• scale (float ) – Scale
• qk_rope_dim (int ) – QK rope dimension
• prev_output (TensorValue | None ) – Optional. Previous output tensor
• prev_softmax_info (TensorValue | None ) – Optional. Previous softmax info tensor

Returns:

• The first tensor is the output tensor for this iteration
• The second tensor is the softmax info tensor for this iteration

Return type:

A tuple of two tensors

flash_attention()

max.nn.kernels.flash_attention(kv_params, input, kv_collection, layer_idx, attention_mask, valid_lengths, scale)

Computes flash attention provided the mo.opaque KV Cache.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection )
• layer_idx (TensorValue )
• attention_mask (TensorValue )
• valid_lengths (TensorValue )
• scale (float )

Return type:

TensorValue

flash_attention_ragged()

max.nn.kernels.flash_attention_ragged(kv_params, input, input_row_offsets, kv_collection, layer_idx, mask_variant, scale, local_window_size=-1)

Computes flash (self) attention provided the !mo.opaque KV Cache.

Notably, this materializes the attention mask (dependent on MHAMaskVariant) within the kernel. input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Note that this is self attention and the KV sequence length is assumed to be equal to the Q sequence length. For KV sequence length != Q sequence length, use cross_attention_ragged.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• layer_idx (TensorValue )
• mask_variant (MHAMaskVariant )
• scale (float )
• local_window_size (int )

Return type:

TensorValue

flash_attention_with_causal_mask()

max.nn.kernels.flash_attention_with_causal_mask(kv_params, input, kv_collection, layer_idx, valid_lengths, scale)

Computes flash attention provided the mo.opaque KV Cache. Notably, materializes the causal mask within the kernel.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection )
• layer_idx (TensorValue )
• valid_lengths (TensorValue )
• scale (float )

Return type:

TensorValue

fused_qk_ragged_rope()

max.nn.kernels.fused_qk_ragged_rope(kv_params, input, input_row_offsets, kv_collection, freqs_cis, layer_idx, interleaved=True)

Computes fused query-key attention with rotary positional encodings and ragged inputs.

Parameters:

• input (TensorValue ) – [batch_size * seq_len, n_heads, head_dim]
• input_row_offsets (TensorValue )
• freqs_cis (TensorValue ) – tensor of shape (max_seq_len * 2, head_dim)
• layer_idx (TensorValue )
• interleaved (bool )
• kv_params (KVCacheParams )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )

Return type:

TensorValue

input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

fused_qk_rope()

max.nn.kernels.fused_qk_rope(kv_params, input, kv_collection, freqs_cis_2d, layer_idx, interleaved=True)

Computes fused query-key attention with rotary positional encodings.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection )
• freqs_cis_2d (TensorValue )
• layer_idx (TensorValue )
• interleaved (bool )

Return type:

TensorValue

fused_qkv_matmul()

max.nn.kernels.fused_qkv_matmul(kv_params, input, wqkv, kv_collection, layer_idx, n_heads)

Computes fused query, key and value projections.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• wqkv (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection )
• layer_idx (TensorValue )
• n_heads (int )

Return type:

TensorValue

fused_qkv_ragged_matmul()

max.nn.kernels.fused_qkv_ragged_matmul(kv_params, input, input_row_offsets, wqkv, kv_collection, layer_idx, n_heads, bias=None)

Computes fused query, key, and value projections with ragged input.

input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Raises:

ValueError – on input shapes/dtypes that are invalid for the kernel.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• wqkv (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• layer_idx (TensorValue )
• n_heads (int )
• bias (TensorValue | None )

Return type:

TensorValue

fused_qkv_ragged_matmul_quantized()

max.nn.kernels.fused_qkv_ragged_matmul_quantized(kv_params, input, input_row_offsets, wqkv, kv_collection, layer_idx, n_heads, quantization_config, perm_idx=None, bias=None)

Computes fused query, key, and value projections with ragged input and quantized weight matrices. A quantization_config must be provided.

input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Raises:

ValueError – on input shapes/dtypes that are invalid for the kernel.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• wqkv (TensorValue )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• layer_idx (TensorValue )
• n_heads (int )
• quantization_config (QuantizationConfig )
• perm_idx (TensorValue | None )
• bias (TensorValue | None )

Return type:

TensorValue

fused_qkv_ragged_matmul_scaled_float8()

max.nn.kernels.fused_qkv_ragged_matmul_scaled_float8(kv_params, input, input_row_offsets, wqkv, kv_collection, layer_idx, n_heads, input_scale, weight_scale, bias=None)

Computes fused query, key, and value projections with ragged input.

input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Raises:

ValueError – on input shapes/dtypes that are invalid for the kernel.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• wqkv (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )
• n_heads (int )
• input_scale (TensorValue )
• weight_scale (TensorValue )
• bias (TensorValue | None )

Return type:

TensorValue

grouped_matmul_ragged()

max.nn.kernels.grouped_matmul_ragged(hidden_states, weight, expert_start_indices, expert_ids, expert_usage_stats_host)

Grouped matmul used in MoE layer.

hidden_states and expert_start_indices are used together to implement the ragged tensor. expert_start_indices indicates where each group starts and ends in hidden_states

expert_ids is the id of the expert for each group in hidden_states

expert_usage_stats_host is the maximum number of tokens assigned to any expert, and the number of active experts.

Parameters:

• hidden_states (TensorValue )
• weight (TensorValue )
• expert_start_indices (TensorValue )
• expert_ids (TensorValue )
• expert_usage_stats_host (TensorValue )

Return type:

TensorValue

kv_cache_get_max_seq_len()

max.nn.kernels.kv_cache_get_max_seq_len(kv_collection)

This kernel returns the maximum sequence length.

Parameters:

kv_collection (PagedKVCacheCollection )

Return type:

TensorValue

matmul_k_cache_ragged()

max.nn.kernels.matmul_k_cache_ragged(kv_params, hidden_states, input_row_offsets, weight, kv_collection, layer_idx)

Computes key projections with ragged input.

hidden_states and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Parameters:

• kv_params (KVCacheParams )
• hidden_states (TensorValue )
• input_row_offsets (TensorValue )
• weight (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )

Return type:

None

matmul_kv_cache_ragged()

max.nn.kernels.matmul_kv_cache_ragged(kv_params, hidden_states, input_row_offsets, weight, kv_collection, layer_idx)

Computes key and value projections with ragged input.

hidden_states and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Parameters:

• kv_params (KVCacheParams )
• hidden_states (TensorValue )
• input_row_offsets (TensorValue )
• weight (TensorValue )
• kv_collection (PagedKVCacheCollection )
• layer_idx (TensorValue )

Return type:

None

matmul_static_scaled_float8()

max.nn.kernels.matmul_static_scaled_float8(input, weight, input_scale, weight_scale)

Parameters:

• input (TensorValue )
• weight (TensorValue )
• input_scale (TensorValue )
• weight_scale (TensorValue )

Return type:

TensorValue

merge_ragged_tensors()

max.nn.kernels.merge_ragged_tensors(a, a_row_offsets, b, b_row_offsets)

Merges two ragged tensors into a single ragged tensor.

Both ragged tensors must have the same batch size (same number of row offsets). This function interleaves the rows from each tensor based on their row offsets.

Parameters:

• a (TensorValue ) – The first ragged tensor of shape [total_a_rows, …].
• a_row_offsets (TensorValue ) – The row offsets of the first ragged tensor,indicating where each batch starts and ends in a.
• b (TensorValue ) – The second ragged tensor of shape [total_b_rows, …].
• b_row_offsets (TensorValue ) – The row offsets of the second ragged tensor, indicating where each batch starts and ends in b.

Returns:

• The merged ragged tensor with shape [total_a_rows + total_b_rows, …].
• The merged row offsets with the same shape as input row offsets.

Return type:

A tuple of two tensors

Example

a = [1, 2, 3, 4, 5, 6] a_row_offsets = [0, 2, 6] b = [7, 8, 9, 10] b_row_offsets = [0, 3, 4]

merged_tensor, merged_row_offsets = merge_ragged_tensors(

a, a_row_offsets, b, b_row_offsets)

merged_tensor = [1, 2, 7, 8, 9, 3, 4, 5, 6, 10] merged_row_offsets = [0, 5, 10]

moe_create_indices()

max.nn.kernels.moe_create_indices(topk_ids, num_local_experts)

Creates indices for the MoE layer.

Parameters:

• topk_ids (TensorValue ) – The expert assignments for each token from the router.
• num_local_experts (int ) – The number of experts on this device.

Returns:

• token_expert_order: The reordered token indices, grouped by assigned expert.
• expert_start_indices: The starting index for each expert’s token group in the reordered sequence.
• restore_token_order: The indices to restore original token ordering after expert computation.
• expert_ids: ids of active experts selected for tokens
• expert_usage_stats: The maximum number of tokens assigned to any expert, and the number of active experts.

Return type:

A tuple of four tensors

null_mask_flash_attention_gpu()

max.nn.kernels.null_mask_flash_attention_gpu(q, k, v, scale)

Computes flash attention using GPU-optimized kernel. :param q: Query tensor of shape [batch, seq_len, num_heads, head_dim] :param k: Key tensor of shape [batch, seq_len, num_heads, head_dim] :param v: Value tensor of shape [batch, seq_len, num_heads, head_dim] :param scale: Scaling factor for attention scores

Parameters:

• q (TensorValue )
• k (TensorValue )
• v (TensorValue )
• scale (float )

Return type:

TensorValue

quantize_dynamic_scaled_float8()

max.nn.kernels.quantize_dynamic_scaled_float8(input, scale_ub=1200.0, group_size_or_per_token=-1, out_type=float8_e4m3fn, scales_type=bfloat16)

Dynamically quantize the input tensor to fp8.

Parameters:

• input (TensorValue ) – The input tensor to quantize.
• scale_ub (float ) – The upper bound of the scale factor.
• group_size_or_per_token (int ) – The group size for quantization. When set to -1, the quantization is column-wise.
• out_type (DType ) – The type of the output tensor.
• scales_type (DType ) – The type of the scales tensor.

Returns:

The quantized tensor and the scales.

Return type:

tuple[TensorValue, TensorValue]

quantize_static_scaled_float8()

max.nn.kernels.quantize_static_scaled_float8(x, scale, scale_is_inverted=True)

Parameters:

• x (TensorValue )
• scale (TensorValue )
• scale_is_inverted (bool )

Return type:

TensorValue

rms_norm_key_cache()

max.nn.kernels.rms_norm_key_cache(kv_params, kv_collection, gamma, epsilon, layer_idx, total_seq_len, input_row_offsets, weight_offset, rms_norm_cols=None)

Computes RMSNorm on the _new_ entries in the KVCache.

This function applies RMSNorm to either all dimensions or a subset of dimensions in each head of the key cache. The size of the gamma tensor determines how many dimensions will be normalized. If gamma’s size doesn’t match head_dim, rms_norm_cols must be explicitly specified to confirm the intention to normalize only a subset of dimensions.

Currently, the KVCacheT class itself isn’t aware of the new cache entries until cache length increment, which happens after model forward. So use input_row_offsets to do this bookkeeping.

Parameters:

• kv_params (KVCacheParams )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• gamma (TensorValue )
• epsilon (float | floating )
• layer_idx (TensorValue )
• total_seq_len (Dim )
• input_row_offsets (TensorValue )
• weight_offset (float | floating )
• rms_norm_cols (int | None )

Return type:

None

swish_glu()

max.nn.kernels.swish_glu(a, b0, b1)

Computes swish(a@b0.t()) * (a@b1.t())

Parameters:

• a (Value [ TensorType ] | TensorValue | Shape | Dim | int | float | integer | floating | ndarray )
• b0 (Value [ TensorType ] | TensorValue | Shape | Dim | int | float | integer | floating | ndarray )
• b1 (Value [ TensorType ] | TensorValue | Shape | Dim | int | float | integer | floating | ndarray )

Return type:

TensorValue

unfused_qkv_ragged_matmul_gguf_quantized()

max.nn.kernels.unfused_qkv_ragged_matmul_gguf_quantized(kv_params, input, input_row_offsets, n_heads, q_weight, k_weight, v_weight, quantization_encoding_q, quantization_encoding_k, quantization_encoding_v, kv_collection, layer_idx)

Computes fused query, key, and value projections with ragged input and quantized weight matrices. A quantization_config must be provided.

input and input_row_offsets are used together to implement the ragged tensor. input_row_offsets indicates where each batch starts and ends in input

Raises:

ValueError – on input shapes/dtypes that are invalid for the kernel.

Parameters:

• kv_params (KVCacheParams )
• input (TensorValue )
• input_row_offsets (TensorValue )
• n_heads (int )
• q_weight (TensorValue )
• k_weight (TensorValue )
• v_weight (TensorValue )
• quantization_encoding_q (QuantizationEncoding )
• quantization_encoding_k (QuantizationEncoding )
• quantization_encoding_v (QuantizationEncoding )
• kv_collection (ContinuousBatchingKVCacheCollection | PagedKVCacheCollection )
• layer_idx (TensorValue )

Return type:

TensorValue

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit