max.nn.kernels

`always_ready_chain`

property always_ready_chain: _ChainValue

source

A graph-global, immutable chain that is always ready.

Created once per graph and never advanced/merged by the graph itself. Use it for operations that are safe to schedule without threading per-device ordering (for example, host→device transfers for staging).

`copy()`

copy()

source

Creates a deep copy of this graph.

The copy shares no MLIR state with the original: staging or lowering on either graph leaves the other untouched. Use this to hand a graph to another thread (for example, background compilation) while continuing to build or execute the original. The kernel library is shared, not copied.

Returns:: A new Graph wrapping a deep copy of this graph’s module.
Raises:: ValueError – If a custom op in this graph has a data-dependent output dim; a copy would restart their per-graph numbering (see _allocate_data_dependent_ordinal()).
Return type:: Graph

`current`

current

source

`debug`

debug = <max.graph.graph.GraphDebugConfig object>

source

`device_chains`

device_chains: _DeviceChainMap

source

`inputs`

property inputs: Sequence[Value[Any]]

source

The input values of the graph.

Returns:: A sequence of Value objects corresponding to the input_types passed at construction, excluding internal chain values.

`kernel_libraries_paths`

property kernel_libraries_paths: list[Path]

source

Returns the list of extra kernel libraries paths for the custom ops.

`module`

property module: Module

source

The Module that owns this graph.

Multiple Graph instances built with the same module= argument share the same underlying Module; that shared object is what you pass to max.engine.InferenceSession.load_all() when compiling several graphs together.

`output()`

output(*outputs)

source

Sets the output values of the graph and finalizes construction.

Call this once after building all ops. The graph can’t be executed until output() has been called. Subsequent calls to output_types read back the types of the values passed here.

For example, build a graph that doubles its input and set the output:

from max.dtype import DType
from max.graph import DeviceRef, Graph, ops
from max.graph.type import TensorType

input_type = TensorType(DType.float32, [4], DeviceRef.CPU())

with Graph("double", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor
    two = ops.constant(2.0, DType.float32, device=DeviceRef.CPU())
    graph.output(ops.elementwise.mul(x, two))

Pass one named tree, such as a dict of logits and hidden states, to name each output by its path in it; output_tree then returns the output types in that structure.

Parameters:: outputs (Value[Any] | Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | Sequence[GraphOutputs] | Mapping[Any, GraphOutputs]) – The output values of the graph, or a single tree whose leaves they are. Each value may be a Value or any TensorValueLike.
Return type:: None

`output_tree`

property output_tree: tree.Tree[Type[Any]]

source

The output types, nested as they were passed to output().

A single named tree, such as a dict of logits and hidden states, names each output by its path in it; separate values, or a graph read back from MLIR, give a tuple.

Raises:: TypeError – If the graph has not yet been terminated by a call to output().

`output_types`

property output_types: list[Type[Any]]

source

The types of the graph output values.

Returns:: A list of Type objects corresponding to the values passed to output(), in the same order.
Raises:: TypeError – If the graph has not yet been terminated by a call to output().

`profile_scope()`

profile_scope(name, color=None)

source

Labels every op created within this block for profiling.

Ops created while the scope is active carry name in their MLIR location, which downstream NVTX tracing surfaces as a [name] suffix on the op’s own per-kernel trace name. Scopes nest: entering a second profile_scope() while the first is still active labels ops with both names, outermost scope first (for example kernel_name [draft_forward/target_forward]).

color is only used by the optional in-region range bracketing mechanism (max-debug.profile-scope-tracing); it never appears on the per-kernel trace name, regardless of what was given.

from max.graph import Graph, ProfileScopeColor
with Graph("main") as graph:
    with graph.profile_scope(
        "draft_forward", color=ProfileScopeColor.ORANGE
    ):
        ...  # ops here trace as "kernel_name [draft_forward]"

Parameters:

name (str) – The scope label to attach to every op created inside this block.
color (ProfileScopeColor | None) – Optional NVTX color for the in-region range bracketing mechanism. Has no effect on the per-kernel trace name.

Return type:

Generator[None]

Note

Scopes are tracked in a single module-level ContextVar. Do not construct a second nested Graph while a sibling graph’s profile_scope is still active; labels would leak between the two graphs.

`InputScaleSpec`

class max.nn.kernels.InputScaleSpec(granularity, origin, dtype, activation_scale_ub=None, block_size=None)

source

Bases: object

Specifies how input activations are scaled for scaled quantization.

Parameters:

granularity (ScaleGranularity)
origin (ScaleOrigin)
dtype (DType)
activation_scale_ub (float | None)
block_size (tuple[int, int] | None)

`activation_scale_ub`

activation_scale_ub: float | None = None

source

An optional upper bound for dynamic activation scaling.

`block_size`

block_size: tuple[int, int] | None = None

source

The tuple[int, int] of the block size for block-wise scaling.

`dtype`

dtype: DType

source

The DType of the input scale factor(s).

`granularity`

granularity: ScaleGranularity

source

The ScaleGranularity of the input scale factor application.

`is_block`

property is_block: bool

source

Whether the input scale granularity is block-wise.

`is_colwise`

property is_colwise: bool

source

Whether the input scale granularity is column-wise.

`is_rowwise`

property is_rowwise: bool

source

Whether the input scale granularity is row-wise.

`is_tensor`

property is_tensor: bool

source

Whether the input scale granularity is per-tensor.

`origin`

origin: ScaleOrigin

source

The ScaleOrigin (static or dynamic) of the input scale factor.

`KVCacheParams`

class max.nn.kernels.KVCacheParams(dtype, head_dim, num_layers, devices, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, per_layer_buffers=False, kv_hash_algo='ahash64', kv_hash_seed=None, kv_connector_config=<factory>, page_size=128, slots_per_page=None, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0, window_size=None, stride_alignment=None, scratch=False)

source

Bases: KVCacheParamInterface

Configuration parameters for key-value cache management in transformer models.

This class encapsulates all configuration options for managing KV caches during inference, including parallelism settings, and memory management.

Parameters:

dtype (DType)
head_dim (int)
num_layers (int)
devices (Sequence[DeviceRef])
enable_prefix_caching (bool)
enable_dp_cross_replica_prefix_copy (bool)
per_layer_buffers (bool)
kv_hash_algo (Literal['ahash64', 'sha256'])
kv_hash_seed (bytes | None)
kv_connector_config (KVConnectorConfigInterface)
page_size (int)
slots_per_page (int | None)
data_parallel_degree (int)
kvcache_quant_config (KVCacheQuantizationConfig | None)
speculative_method (Literal['eagle', 'mtp', 'dflash', 'dflash2', 'dspark'] | None)
num_draft_tokens (int)
window_size (int | None)
stride_alignment (int | None)
scratch (bool)

`allocate_buffers()`

allocate_buffers(total_num_pages, _prefix='')

source

Allocates the buffers for the KV cache.

Parameters:

total_num_pages (int)
_prefix (str)

Return type:

list[KVCacheBuffer]

`build_runtime_inputs()`

build_runtime_inputs(assignments, buffers, _prefix='', *, staging=None)

source

Builds the runtime KV-cache leaf spanning all replicas.

assignments and buffers are indexed by data-parallel replica. The returned KVCacheInputs lists one KVCacheInputsPerDevice per (replica, TP shard), in the same replica-major order as get_symbolic_inputs().

Parameters:

assignments (Sequence[KVCacheAssignments])
buffers (Sequence[KVCacheBufferInterface])
_prefix (str)
staging (GraphInputStagingInterface | None)

Return type:

tuple[KVCacheInputsPerDevice[Buffer, Buffer], …]

`bytes_per_block`

property bytes_per_block: int

source

Returns the number of bytes per cache block.

When TP>1, each block is sharded across the devices in the tensor parallel group. This method returns the total memory needed to store a block across these devices. Includes memory needed for scales if quantization is enabled.

Returns:: The number of bytes per cache block.

`bytes_per_scale_block`

property bytes_per_scale_block: int

source

Returns the number of bytes per scale block.

`bytes_per_value_block`

property bytes_per_value_block: int

source

Returns the number of bytes per value block.

`data_parallel_degree`

data_parallel_degree: int = 1

source

Degree of data parallelism. Devices are grouped replica-major, with n_devices // data_parallel_degree TP shards per replica.

`devices`

devices: Sequence[DeviceRef]

source

Devices to use for the KV cache.

`devices_per_replica`

property devices_per_replica: Sequence[Sequence[DeviceRef]]

source

Returns the devices per replica.

`dtype`

dtype: DType

source

Data type for storing key and value tensors in the cache.

`dtype_shorthand`

property dtype_shorthand: str

source

Returns a shorthand textual representation of the data type.

Returns:: “bf16” for bfloat16 dtype, “f32” otherwise.

`enable_dp_cross_replica_prefix_copy`

enable_dp_cross_replica_prefix_copy: bool = True

source

Whether a prefix-cache block resident on another data-parallel (DP) replica’s device may be materialized locally via a device-to-device copy to serve a cache hit. When False, cross-replica reuse is only served from the shared external tier via the KV connector (or recomputed). Only relevant when data_parallel_degree > 1 and prefix caching is enabled.

`enable_prefix_caching`

enable_prefix_caching: bool = False

source

Whether to enable prefix caching for efficient reuse of common prompt prefixes.

`get_symbolic_inputs()`

get_symbolic_inputs(namespace='')

source

Computes the symbolic inputs for the KV cache.

Parameters:: namespace (str) – Prefix disambiguating this cache’s per-pool page-count dim from sibling caches in a multi-group tree (empty for a single-group cache).
Returns:: The symbolic inputs for the KV cache.
Return type:: tuple[KVCacheInputsPerDevice[TensorType, BufferType], …]

`group_id`

property group_id: KVCacheGroupId

source

Returns the group id this cache pools under: scratch for a scratch cache, sliding_window with the window size when one is set, else full.

`head_dim`

head_dim: int

source

Dimensionality of each attention head.

`is_fp8_kv_dtype`

property is_fp8_kv_dtype: bool

source

Whether the KV cache stores FP8 data, for dispatch resolution.

Unlike quantized_kv_cache (which also requires valid scale config), this checks only the storage dtype—matching the compile-time detection in the MLA decode kernel.

TODO(SERVOPT-1094): Once SnapMLA uses a valid scale_dtype, this can be replaced by quantized_kv_cache.

`is_nvfp4_kv_cache`

property is_nvfp4_kv_cache: bool

source

Whether the cache stores packed NVFP4: uint8 values holding two E2M1 elements each, with E4M3 scales.

`kv_cache_scale_dtype`

property kv_cache_scale_dtype: DType

source

Returns the dtype of the KV cache scales.

Returns:: The dtype of the KV cache scales.

`kv_connector_config`

kv_connector_config: KVConnectorConfigInterface

source

Holds the connector type and its settings. The default is a null connector (no external caching).

`kv_dim`

property kv_dim: int

source

Returns the number of key/value tensors each cache slot holds.

`kv_hash_algo`

kv_hash_algo: Literal['ahash64', 'sha256'] = 'ahash64'

source

Hash algorithm used for KV-cache block identity.

`kv_hash_seed`

kv_hash_seed: bytes | None = None

source

Resolved 32-byte cluster seed for sha256. None for ahash64.

Set by KVCacheConfig.to_params via resolve_kv_hash_seed.

`kvcache_quant_config`

kvcache_quant_config: KVCacheQuantizationConfig | None = None

source

KVCache quantization config. Currently only FP8 quantization supported.

`leaf_id()`

leaf_id(_prefix='')

source

Returns the id this cache’s pages are pooled under.

Parameters:: _prefix (str)
Return type:: str

`leaf_kind`

leaf_kind: ClassVar[CacheLeafKind] = 'attention'

source

`leaves()`

leaves(_prefix='')

source

Returns the leaves of the KV cache.

A leaf reports what one device holds, the unit a slab tiles in, where bytes_per_block counts a block across the whole replica. The division is exact: that figure is this one times the degree.

Parameters:: _prefix (str)
Return type:: Mapping[str, KVLeafRegion]

`n_devices`

property n_devices: int

source

Returns the number of devices.

Returns:: The number of devices.

`n_kv_heads_per_device`

property n_kv_heads_per_device: int

source

Returns the number of KV attention heads on one device.

`num_draft_tokens`

num_draft_tokens: int = 0

source

Total draft tokens generated per speculative iteration.

Zero when no speculative decoding is configured.

`num_layers`

num_layers: int

source

Number of layers in the model.

`page_size`

page_size: int = 128

source

Number of tokens per page (block).

This value is expressed in tokens, not bytes. The byte footprint of a page is derived from pipeline configuration.

Current constraints: the page size must be a multiple of 128 and at least 128.

`per_layer_buffers`

per_layer_buffers: bool = False

source

When True, allocate one standalone single-layer buffer per layer instead of one [..., num_layers, ...] multi-layer buffer.

Each attention dispatch then binds only its own per-layer buffer, so the pool total can exceed a per-allocation size cap (such as a device’s maximum single allocation) while every individual buffer stays under it. Defaults to False (one multi-layer buffer), keeping all other backends and models byte-identical.

`quantized_kv_cache`

property quantized_kv_cache: bool

source

Returns whether KV cache quantization is enabled.

`replicates_kv_across_tp`

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

`row_bytes`

property row_bytes: int

source

Returns one value row, num_heads * head_size * dtype_size.

`scale_row_bytes`

property scale_row_bytes: int

source

Returns one scale row; see row_bytes.

`scratch`

scratch: bool = False

source

Whether the cache is one block a request, kept until release and never published. Excludes window_size.

`shape_per_block`

property shape_per_block: list[int]

source

Returns the shape of each cache block.

The slot dimension is slots_per_page, which equals page_size unless the architecture stores a page’s tokens compressed.

Returns:: The shape of the cache block.

`shape_per_layer_block`

property shape_per_layer_block: list[int]

source

Returns the block shape for a single-layer buffer.

Same as shape_per_block but with the layer dimension pinned to 1. Used when per_layer_buffers is set: the pool allocates num_layers such buffers per device instead of one multi-layer buffer. The attention kernel derives num_layers from this dim, so a single-layer buffer (num_layers == 1) with layer_idx == 0 is self-consistent.

`shape_per_layer_scale_block`

property shape_per_layer_scale_block: list[int]

source

Scale-block shape for a single-layer buffer (layer dim pinned to 1).

The scale analog of shape_per_layer_block: used with per_layer_buffers on a quantized cache, where the pool allocates one single-layer scale buffer per layer instead of one multi-layer one.

`shape_per_scale_block`

property shape_per_scale_block: list[int]

source

Returns the shape of each scale block used for KVCache quantization

Returns:: The shape of the KVCache quantization scales block.

`slab_to_bound_views()`

slab_to_bound_views(slabs)

source

Returns nothing: these pages are addressed by the lookup table.

Parameters:: slabs (Sequence[Buffer])
Return type:: Mapping[str, list[Buffer]]

`slab_to_buffer_views()`

slab_to_buffer_views(buffers, padded_page_bytes=None, _prefix='')

source

Converts a slab of memory into per-leaf page views.

Each view carries its own page-to-page distance in its dim-0 stride, so nothing downstream – the pool, the connector, the transfer engine – has to be told whether this pool padded its pages.

Parameters:

buffers (Sequence[Buffer]) – One raw arena slab per TP shard.
padded_page_bytes (Mapping[str, int] | None) – The pool’s padded page size per leaf, keyed as leaves() keys it. None means nothing was padded.
_prefix (str) – Leaf-id prefix identifying this node in a params tree.

Returns:

The buffer views for this cache.

Return type:

KVCacheBufferInterface

`slots_per_page`

slots_per_page: int | None = None

source

Number of storage slots a page holds, which may be fewer than the page_size tokens it covers.

None (the default) means one slot per token; construction resolves it to page_size, so every later reader sees an int. An architecture whose attention compresses a run of tokens into a single cached entry – a pooled key, a compressed latent – sets this to the number of entries a page’s tokens compress into, and gets a page that many slots deep instead. It must divide page_size evenly, so a page never ends mid-slot.

Only the physical buffer shape and the kernels reading it are affected: block allocation, the page table and prefix caching stay in tokens, and a block still covers page_size of them.

`speculative_method`

speculative_method: Literal['eagle', 'mtp', 'dflash', 'dflash2', 'dspark'] | None = None

source

Speculative decoding method propagated from SpeculativeConfig.

`storage_head_dim`

property storage_head_dim: int

source

Returns the stored width of one head row in dtype elements.

Equals head_dim except for packed NVFP4, which stores two elements per byte, apart from its trailing FP8 elements (see KVCacheQuantizationConfig.fp8_tail_dim).

`stride_alignment`

stride_alignment: int | None = None

source

The page stride alignment this cache’s leaves declare, or None to keep them packed. See KVLeafRegion.stride_alignment.

`tensor_parallel_degree`

property tensor_parallel_degree: int

source

Returns the tensor parallel degree.

Returns:: The tensor parallel degree.

`unflatten_basic_kv_tree()`

unflatten_basic_kv_tree(it)

source

Unflattens a basic KV tree from a graph-input iterator.

Requires that the model is a basic height-1 tree. This method does not work on nested trees.

Parameters:: it (Iterator[Any])
Return type:: tuple[list[KVCacheInputsPerDevice[TensorValue, BufferValue]], …]

`unflatten_kv_inputs()`

unflatten_kv_inputs(it)

source

Unflattens the KV cache inputs from a graph-input iterator.

Parameters:: it (Iterator[Any])
Return type:: tuple[KVCacheInputsPerDevice[TensorValue, BufferValue], …]

`window_size`

window_size: int | None = None

source

Window size for the sliding window attention. None for global attention.

`MHAKVCacheParams`

class max.nn.kernels.MHAKVCacheParams(dtype, head_dim, num_layers, devices, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, per_layer_buffers=False, kv_hash_algo='ahash64', kv_hash_seed=None, kv_connector_config=<factory>, page_size=128, slots_per_page=None, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0, window_size=None, stride_alignment=None, scratch=False, *, n_kv_heads, allow_kv_head_replication=False)

source

Bases: KVCacheParams

KV cache parameters for multi-head attention (MHA).

Parameters:

dtype (DType)
head_dim (int)
num_layers (int)
devices (Sequence[DeviceRef])
enable_prefix_caching (bool)
enable_dp_cross_replica_prefix_copy (bool)
per_layer_buffers (bool)
kv_hash_algo (Literal['ahash64', 'sha256'])
kv_hash_seed (bytes | None)
kv_connector_config (KVConnectorConfigInterface)
page_size (int)
slots_per_page (int | None)
data_parallel_degree (int)
kvcache_quant_config (KVCacheQuantizationConfig | None)
speculative_method (Literal['eagle', 'mtp', 'dflash', 'dflash2', 'dspark'] | None)
num_draft_tokens (int)
window_size (int | None)
stride_alignment (int | None)
scratch (bool)
n_kv_heads (int)
allow_kv_head_replication (bool)

`allow_kv_head_replication`

allow_kv_head_replication: bool = False

source

Allows TP wider than n_kv_heads. When set and n_devices is a multiple of n_kv_heads, each KV head is replicated across a group of devices (n_kv_heads_per_device == 1).

`devices`

devices: Sequence[DeviceRef]

source

Devices to use for the KV cache.

`dispatch_metadata_spec()`

dispatch_metadata_spec()

source

Returns the spec MHAAttnKey packs into.

Return type:: DispatchMetadataSpec

`dtype`

dtype: DType

source

Data type for storing key and value tensors in the cache.

`graph_capture_probe_cache_lengths()`

graph_capture_probe_cache_lengths(max_cache_length, q_max_seq_len=1)

source

Returns cache lengths to probe for distinct num_partitions.

Parameters:

max_cache_length (int)
q_max_seq_len (int)

Return type:

list[int]

`head_dim`

head_dim: int

source

Dimensionality of each attention head.

`kv_connector_config`

kv_connector_config: KVConnectorConfigInterface

source

Holds the connector type and its settings. The default is a null connector (no external caching).

`kv_dim`

property kv_dim: int

source

Returns two: each slot holds a key and a value tensor.

`n_kv_heads`

n_kv_heads: int

source

Total number of key-value attention heads across all devices.

`n_kv_heads_per_device`

property n_kv_heads_per_device: int

source

Returns the KV heads on one device, or 1 per device group when heads are replicated (allow_kv_head_replication).

`num_layers`

num_layers: int

source

Number of layers in the model.

`replicates_kv_across_tp`

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

`resolve_attn_key()`

resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)

source

Resolves the decode attention dispatch shape for the given shape.

Parameters:

batch_size (int) – Number of requests in the decode batch.
max_prompt_length (int) – Per-step query width (1 for plain decode, 1 + num_spec_tokens for speculative verify).
max_cache_valid_length (int) – Maximum valid cache length in the batch.

Returns:

The resolved AttnKeyInterface.

Return type:

AttnKeyInterface

`MHAMaskVariant`

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

source

Bases: str, Enum

Defines the integer mask variant codes used by multihead attention kernels.

`CAUSAL_MASK`

CAUSAL_MASK = '0'

source

Causal masking; each position attends only to itself and earlier positions.

`CHUNKED_CAUSAL_MASK`

CHUNKED_CAUSAL_MASK = '3'

source

Causal masking applied over fixed-size chunks of the sequence.

`NULL_MASK`

NULL_MASK = '2'

source

No masking; every position attends to every position.

`SLIDING_WINDOW_CAUSAL_MASK`

SLIDING_WINDOW_CAUSAL_MASK = '4'

source

Each position attends to itself and the previous window_size - 1 positions.

`SLIDING_WINDOW_NONCAUSAL_MASK`

SLIDING_WINDOW_NONCAUSAL_MASK = '5'

source

Each position attends to itself, the previous window_size - 1 positions, and every later position.

`MXFP6_LANE_BYTES`

max.nn.kernels.MXFP6_LANE_BYTES = 24

source

Bytes one lane feeds the MFMA for FP6 (32 elements x 6 bits).

MXFP4 and MXFP8 both use 16-byte lanes (FP8 splits its 32-byte fragment into two 16-byte halves at a K stride), so 24 is the only width that needs the plane-split layout.

`NVFP4_MLA_PREFILL_CONVERT_MAX_ROWS`

max.nn.kernels.NVFP4_MLA_PREFILL_CONVERT_MAX_ROWS = 524288

source

Default most context rows that conversion converts (576 bytes each). Past it the NVFP4 kernel runs. Sized so a single request’s prefill converts up to a 512k-token context, where converting stays well ahead of the NVFP4 kernel. Override with MODULAR_MLA_NVFP4_PREFILL_CONVERT_MAX_ROWS.

`NVFP4_MLA_PREFILL_CONVERT_RATIO`

max.nn.kernels.NVFP4_MLA_PREFILL_CONVERT_RATIO = 150

source

Default most context rows per query token at which sparse MLA prefill over an NVFP4 cache converts the context to FP8 once instead of expanding each selected row per query token. Override with MODULAR_MLA_NVFP4_PREFILL_CONVERT_RATIO; 0 turns the conversion off.

`QuantConfig`

class max.nn.kernels.QuantConfig(input_scale, weight_scale, mlp_quantized_layers, attn_quantized_layers, format, embedding_output_dtype=None, shared_experts_weight_dtype=None, bias_dtype=None, can_use_fused_mlp=False, can_use_fused_swiglu=False, scales_pre_interleaved=False, _mxfp6_element_format='e2m3', source=WeightSource.CHECKPOINT)

source

Bases: object

Configures scaled quantization settings for a layer or model section.

For example, to configure NVFP4 block-scaled quantization for all layers in a 19-layer model:

from max.dtype import DType
from max.nn import QuantConfig, QuantFormat
from max.nn.quant_config import (
    InputScaleSpec,
    ScaleGranularity,
    ScaleOrigin,
    WeightScaleSpec,
)

all_layers = set(range(19))

input_spec = InputScaleSpec(
    granularity=ScaleGranularity.BLOCK,
    origin=ScaleOrigin.STATIC,
    dtype=DType.float32,
    block_size=(1, 16),
)
weight_spec = WeightScaleSpec(
    granularity=ScaleGranularity.BLOCK,
    dtype=DType.float8_e4m3fn,
    block_size=(1, 8),
)
config = QuantConfig(
    input_scale=input_spec,
    weight_scale=weight_spec,
    mlp_quantized_layers=all_layers,
    attn_quantized_layers=all_layers,
    format=QuantFormat.NVFP4,
)

Parameters:

input_scale (InputScaleSpec)
weight_scale (WeightScaleSpec)
mlp_quantized_layers (set[int])
attn_quantized_layers (set[int])
format (QuantFormat)
embedding_output_dtype (DType | None)
shared_experts_weight_dtype (DType | None)
bias_dtype (DType | None)
can_use_fused_mlp (bool)
can_use_fused_swiglu (bool)
scales_pre_interleaved (bool)
_mxfp6_element_format (str)
source (WeightSource)

`attn_quantized_layers`

attn_quantized_layers: set[int]

source

Set of layer indices with quantized attention projections.

Attention projections are quantized on an all-or-nothing basis per layer: either all of q_proj, k_proj, v_proj, and o_proj are quantized, or all four remain in bfloat16.

`bias_dtype`

bias_dtype: DType | None = None

source

The DType of bias weights.

`blockscaled_fp8()`

classmethod blockscaled_fp8(*, mlp_quantized_layers, attn_quantized_layers, block_size=128, scale_dtype=float32, embedding_output_dtype=None, bias_dtype=None)

source

Builds a block-scaled FP8 E4M3 config, as in DeepSeek-V3’s fp8.

Each weight has one scale per block_size x block_size block. Activations are quantized at runtime with one scale per token per block_size inputs.

Parameters:

mlp_quantized_layers (Iterable[int]) – Layer indices whose MLPs are quantized.
attn_quantized_layers (Iterable[int]) – Layer indices whose attention projections are quantized.
block_size (int) – The side of a weight scale block, and the width of an activation scale block.
scale_dtype (DType) – The dtype of the weight and activation scales.
embedding_output_dtype (DType | None) – The dtype of the embedding output.
bias_dtype (DType | None) – The dtype of the bias weights.

Returns:

The config.

Return type:

`can_use_fused_mlp`

can_use_fused_mlp: bool = False

source

Whether the quantization scales can be used with fused MLP operations.

`can_use_fused_swiglu`

can_use_fused_swiglu: bool = False

source

Whether to use the fused grouped matmul + SwiGLU + NVFP4/MXFP4/MXFP8 quant SM100 kernel for the MoE gate/up projection. When True, the MoE layer pre-permutes gate_up_proj and its scales on the N axis (sigma(2i)=i, sigma(2i+1)=D+i) and dispatches the internal grouped_matmul_blocked_swiglu kernel wrapper. Defaults to False so the chained (matmul -> BF16 -> SwiGLU+quant) path is unchanged.

`embedding_output_dtype`

embedding_output_dtype: DType | None = None

source

The DType of the output from the embedding layer.

`format`

format: QuantFormat

source

The QuantFormat identifying the quantization format.

`input_scale`

input_scale: InputScaleSpec

source

InputScaleSpec for input activation scaling.

`is_dynamic`

property is_dynamic: bool

source

True if this input scale is dynamic.

`is_fp4`

property is_fp4: bool

source

True if this config represents any FP4 variant (NVFP4 or MXFP4).

`is_int8_w8a8`

property is_int8_w8a8: bool

source

True if this config represents symmetric int8 W8A8.

`is_mx`

property is_mx: bool

source

True for any OCP microscaling format (MXFP4, MXFP6 or MXFP8).

These share the E8M0 per-32-element block scale; NVFP4 is excluded, as it scales per 16 elements with an FP8 scale.

`is_mxfp4`

property is_mxfp4: bool

source

Returns True if this config represents MXFP4 quantization.

`is_mxfp6`

property is_mxfp6: bool

source

Returns True if this config represents MXFP6 quantization.

`is_mxfp8`

property is_mxfp8: bool

source

Returns True if this config represents MXFP8 quantization.

`is_nvfp4`

property is_nvfp4: bool

source

True if this config represents modelopt NVFP4.

`is_static`

property is_static: bool

source

True if this input scale is static.

`mlp_quantized_layers`

mlp_quantized_layers: set[int]

source

Set of layer indices with quantized MLPs.

MLPs are quantized on an all-or-nothing basis per layer: either all of gate_proj, down_proj, and up_proj are quantized, or all three remain in bfloat16.

`mx_element_dtype`

property mx_element_dtype: DType

source

The dtype of one quantized element, for MX formats.

This is the logical element type, not the storage type: sub-byte MX weights are stored as packed uint8 bytes, and that spelling cannot say which format the bytes hold – MXFP4 and MXFP6 share it. Code that needs to branch on the format should branch on this instead of pairing dtype == DType.uint8 with an is_mxfp6 flag.

Raises:: ValueError – If this config is not an MX format.

`mxfp4()`

classmethod mxfp4(*, mlp_quantized_layers, attn_quantized_layers=(), embedding_output_dtype=None, bias_dtype=None, can_use_fused_mlp=False)

source

Builds an OCP MXFP4 config with dynamic MXFP8 activations.

Weights are packed E2M1 codes with one E8M0 scale per 32 inputs. Activations are quantized at runtime over the same 32-input blocks.

Parameters:

mlp_quantized_layers (Iterable[int]) – Layer indices whose MLPs are quantized.
attn_quantized_layers (Iterable[int]) – Layer indices whose attention projections are quantized.
embedding_output_dtype (DType | None) – The dtype of the embedding output.
bias_dtype (DType | None) – The dtype of the bias weights.
can_use_fused_mlp (bool) – Whether the fused MLP path may be used.

Returns:

The config.

Return type:

`mxfp6_format`

property mxfp6_format: str

source

The FP6 element encoding, "e2m3" or "e3m2".

E2M3 is the default: with 3 mantissa bits against E3M2’s 2, it is the better choice for weights, whose dynamic range a per-32 block scale already absorbs.

`mxfp8()`

classmethod mxfp8(*, mlp_quantized_layers, attn_quantized_layers=(), embedding_output_dtype=None, bias_dtype=None, can_use_fused_mlp=False)

source

Builds an OCP MXFP8 config with dynamic MXFP8 activations.

Weights are unpacked E4M3 values with one E8M0 scale per 32 inputs. Unlike NVFP4 there is no per-tensor second scale level: MXFP8 is pure block scaling, so nothing rescales the matmul output.

Parameters:

mlp_quantized_layers (Iterable[int]) – Layer indices whose MLPs are quantized.
attn_quantized_layers (Iterable[int]) – Layer indices whose attention projections are quantized.
embedding_output_dtype (DType | None) – The dtype of the embedding output.
bias_dtype (DType | None) – The dtype of the bias weights.
can_use_fused_mlp (bool) – Whether the fused MLP path may be used.

Returns:

The config.

Return type:

`nvfp4()`

classmethod nvfp4(*, mlp_quantized_layers, attn_quantized_layers, embedding_output_dtype=None, bias_dtype=None, shared_experts_weight_dtype=None, can_use_fused_mlp=False)

source

Builds a modelopt NVFP4 config.

Weights are packed E2M1 codes with one E4M3 scale per 16 inputs and a float32 per-tensor scale. The activation’s per-tensor scale is static, calibrated into the checkpoint; its 16-input block scales are computed at runtime against it.

Parameters:

mlp_quantized_layers (Iterable[int]) – Layer indices whose MLPs are quantized.
attn_quantized_layers (Iterable[int]) – Layer indices whose attention projections are quantized.
embedding_output_dtype (DType | None) – The dtype of the embedding output.
bias_dtype (DType | None) – The dtype of the bias weights.
shared_experts_weight_dtype (DType | None) – The weight dtype of the MoE shared experts, when it differs from the routed experts’.
can_use_fused_mlp (bool) – Whether the fused MLP path may be used.

Returns:

The config.

Return type:

`quantized_scales_type()`

quantized_scales_type(quantized_shape, device_ref)

source

The TensorType of the scales tensor after dynamic quantization.

Parameters:

quantized_shape (Shape)
device_ref (DeviceRef)

Return type:

TensorType

`quantizes_on_load`

property quantizes_on_load: bool

source

Whether weights arrive in native precision and are quantized here.

The format says what weights become; this says whether the checkpoint already holds them that way.

`scales_granularity_mnk`

property scales_granularity_mnk: tuple[int, int, int]

source

The weight and input scale granularities on the M, N, and K axes.

`scales_pre_interleaved`

scales_pre_interleaved: bool = False

source

Whether weight scales in the checkpoint are already stored in the 5D TCGEN-interleaved layout expected by the FP4 matmul kernel (NVFP4 only). Note that scales in the 5D TCGEN-interleaved layout are typically flattened to 2D [M, K//16] in the checkpoint.

`shared_experts_dtype()`

shared_experts_dtype(routed_weight_dtype)

source

Resolve weight dtype for MoE shared-expert MLPs.

Parameters:: routed_weight_dtype (DType)
Return type:: DType

`shared_experts_use_quant()`

shared_experts_use_quant(routed_weight_dtype)

source

Whether shared experts use the same quantized weights as routed experts.

Parameters:: routed_weight_dtype (DType)
Return type:: bool

`shared_experts_weight_dtype`

shared_experts_weight_dtype: DType | None = None

source

Weight storage dtype for MoE shared-expert MLPs when they differ from routed experts.

When None, shared experts use the same dtype and quantization as routed experts. When set (such as bfloat16 for mixed Kimi K2.6 NVFP4 checkpoints), shared-expert linears omit quant_config while routed experts remain quantized.

`source`

source: WeightSource = 'checkpoint'

source

Where the weight’s bytes come from. See WeightSource.

`weight_scale`

weight_scale: WeightScaleSpec

source

WeightScaleSpec for weight scaling.

`QuantizationConfig`

class max.nn.kernels.QuantizationConfig(quant_method, bits, group_size, desc_act=False, sym=False)

source

Bases: object

Configuration for specifying quantization parameters that affect inference.

These parameters control how tensor values are quantized, including the method, bit precision, grouping, and other characteristics that affect the trade-off between model size, inference speed, and accuracy.

Parameters:

quant_method (str)
bits (int)
group_size (int)
desc_act (bool)
sym (bool)

`bits`

bits: int

source

The number of bits used to represent each quantized weight element.

`desc_act`

desc_act: bool = False

source

Whether to use activation ordering (descending activation order). Defaults to False.

`group_size`

group_size: int

source

The number of weight elements that share a single set of quantization parameters.

`quant_method`

quant_method: str

source

The quantization method name (for example, gptq or awq).

`sym`

sym: bool = False

source

Whether to use symmetric quantization. Defaults to False.

`QuantizationEncoding`

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

source

Bases: Enum

Quantization encodings supported by MAX Graph.

Quantization reduces the precision of neural network weights to decrease memory usage and potentially improve inference speed. Each encoding represents a different compression method with specific trade-offs between model size, accuracy, and computational efficiency. These encodings are commonly used with pre-quantized model checkpoints (especially GGUF format) or can be applied during weight allocation.

The following example shows how to create a quantized weight using the Q4_K encoding:

from max.dtype import DType
from max.graph import DeviceRef, Weight
from max.graph.quantization import QuantizationEncoding

# Create a quantized weight using Q4_K encoding
encoding = QuantizationEncoding.Q4_K
quantized_weight = Weight(
    name="linear.weight",
    dtype=DType.uint8,
    shape=[4096, 4096],
    device=DeviceRef.GPU(0),
    quantization_encoding=encoding
)

MAX supports several quantization formats optimized for different use cases.

`GPTQ`

GPTQ = 'GPTQ'

source

Group-wise Post-Training Quantization for large language models.

`Q4_0`

Q4_0 = 'Q4_0'

source

Basic 4-bit quantization with 32 elements per block.

`Q4_K`

Q4_K = 'Q4_K'

source

4-bit K-quantization with 256 elements per block.

`Q5_K`

Q5_K = 'Q5_K'

source

5-bit K-quantization with 256 elements per block.

`Q6_K`

Q6_K = 'Q6_K'

source

6-bit K-quantization with 256 elements per block.

`block_parameters`

property block_parameters: BlockParameters

source

Gets the block parameters for this quantization encoding.

Returns:: The parameters describing how elements are organized and encoded in blocks for this quantization encoding.
Return type:: BlockParameters

`block_size`

property block_size: int

source

Number of bytes in encoded representation of block.

All quantization types currently supported by MAX Graph are block-based: groups of a fixed number of elements are formed, and each group is quantized together into a fixed-size output block. This value is the number of bytes resulting after encoding a single block.

Returns:: Size in bytes of each encoded quantization block.
Return type:: int

`elements_per_block`

property elements_per_block: int

source

Number of elements per block.

Returns:: Number of original tensor elements in each quantized block.
Return type:: int

`is_gguf`

property is_gguf: bool

source

Checks if this quantization encoding is compatible with GGUF format.

GGUF is a format for storing large language models and compatible quantized weights.

Returns:: True if this encoding is compatible with GGUF, False otherwise.
Return type:: bool

`name`

property name: str

source

Gets the lowercase name of the quantization encoding.

Returns:: Lowercase string representation of the quantization encoding.
Return type:: str

`StaticDim`

class max.nn.kernels.StaticDim(value)

source

Bases: Dim

A static tensor dimension with a fixed size.

Because a static dimension’s size is fixed, related computation can be optimized at compile time. This is key to good model performance.

The following example creates static dimensions implicitly by passing integer values to TensorType:

from max.dtype import DType
from max.graph import DeviceRef, TensorType
tensor = TensorType(DType.int64, (4, 5), device=DeviceRef.CPU())
# This creates a tensor with 2 static dimensions: 4 and 5 respectively

Converts valid input values to Dim.

Parameters:: dim (int)

`dim`

dim: int

source

The size of the static dimension.

`from_mlir()`

static from_mlir(attr)

source

Constructs a StaticDim from a builtin.IntegerAttr.

Parameters:: attr (TypedAttr) – The builtin.IntegerAttr to parse into a StaticDim.
Returns:: The StaticDim represented by the builtin.IntegerAttr.
Return type:: StaticDim

`parameters`

property parameters: Iterable[SymbolicDim]

source

Lists the symbolic dimension names on which this dim depends.

`substitute()`

substitute(mapping)

source

Returns this dim unchanged: a static dim has no symbols.

Parameters:: mapping (Mapping[str, int | str | Dim | integer | TypedAttr]) – A mapping from symbolic dimension name to the replacement dim (or dim-like value). Ignored.
Returns:: This dim, unchanged.
Return type:: Dim

`to_mlir()`

to_mlir()

source

Creates an mlir.Attribute representing this dimension.

This is used internally when constructing tensor MLIR types.

Returns:: An mlir.Attribute in the context representing the dimension.
Return type:: TypedAttr

`TensorType`

class max.nn.kernels.TensorType(dtype, shape, device, _layout=None)

source

Bases: _TensorTypeBase[TensorType]

A symbolic tensor type.

Use TensorType to declare the expected dtype, shape, and target device of tensor values that flow through a graph during model execution. Unlike an eager tensor, a TensorType holds no data. It is a purely symbolic description of a value’s type at a specific point in the computation. The graph compiler uses this information for shape inference and optimization during graph construction.

The following example shows how to create a tensor type and access its properties:

from max.graph import TensorType, DeviceRef
from max.dtype import DType
# Create a tensor type with float32 elements and static dimensions 2x3
tensor_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
print(tensor_type.dtype)  # Outputs: DType.float32
print(tensor_type.shape)  # Outputs: [2, 3]

A shape’s dimensions can be static (integers), symbolic (strings), or algebraic (expressions over symbolic dimensions). In each case the rank is known at graph construction time.

Pass TensorType instances to load() or Module.compile() (experimental) to define the input types of a graph or model.

Parameters:

dtype (DType) – The data type of the tensor elements.
shape (Shape) – The shape of the tensor, expressed as a Shape.
device (DeviceRef) – The device the tensor is located on. Use DeviceRef.CPU() or DeviceRef.GPU() to create a device reference.
_layout (FilterLayout | None)

`as_buffer()`

as_buffer()

source

Returns the analogous buffer type.

Return type:: BufferType

`from_mlir()`

classmethod from_mlir(type)

source

Constructs a tensor type from an MLIR type.

Parameters:: type (TensorType) – The MLIR Type to parse into a tensor type.
Returns:: The tensor type represented by the MLIR Type value.
Return type:: TensorType

`to_mlir()`

to_mlir()

source

Converts to an mlir.Type instance.

Returns:: An mlir.Type in the specified context.
Return type:: TensorType

`TensorValue`

class max.nn.kernels.TensorValue(value)

source

Bases: Value[TensorType]

Represents a value semantic tensor within a Graph.

It provides various methods and properties to manipulate and query tensor attributes such as shape, data type (dtype), device placement (device), and more.

The following example demonstrates how to create and manipulate tensor values in a graph:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a sample matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("tensor_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Access tensor properties
    print(f"Shape: {tensor.shape}")  # Output: [2, 2]
    print(f"Data type: {tensor.dtype}")  # Output: DType.float32

    # Perform operations on the tensor
    transposed = tensor.T
    doubled = tensor * 2

    print(f"Original shape: {tensor.shape}")  # Output: [2, 2]
    print(f"Transposed shape: {transposed.shape}")  # Output: [2, 2]

Initializes a TensorValue from a tensor-like value.

Parameters:: value (TensorValueLike) – The value to wrap. Can be an MLIR tensor value, another TensorValue, a Dim, or a Shape.

`T`

property T: TensorValue

source

Returns the transposed tensor.

T is the shorthand notation for transposing. For more information, see transpose().

Returns:: A new TensorValue with swapped dimensions.

`argmax()`

argmax(axis=-1)

source

Reduces the tensor using an argmax operation along axis.

When the result is ambiguous ie. there are multiple maxima, selects one index arbitrarily.

from max.dtype import DType
from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("argmax_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Argmax along axis 1 (last dimension of each row)
    indices = x.argmax(axis=1)

    print(f"Input shape: {x.shape}")       # [2, 3]
    print(f"Argmax shape: {indices.shape}")  # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue of dtype DType.int64 with the same rank as the input, and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`broadcast_to()`

broadcast_to(shape)

source

Broadcasts the tensor to a new shape.

The following example demonstrates how to broadcast a tensor to a larger shape:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x2 matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("broadcast_to_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Broadcast tensor to a 3x2x2 tensor (add a new dimension of size 3)
    broadcasted_tensor = tensor.broadcast_to((3, 2, 2))

    print(f"Original shape: {tensor.shape}")  # Output: [2, 2]
    print(f"Broadcasted shape: {broadcasted_tensor.shape}")  # Output: [3, 2, 2]

Parameters:: shape (Iterable[int | str | Dim | integer | TypedAttr]) – An iterable of integers or symbolic dimensions.
Returns:: A new TensorValue with the broadcasted shape.
Return type:: TensorValue

`cast()`

cast(dtype)

source

Casts a symbolic tensor to a different data type.

The following example demonstrates how to cast a tensor from one data type to another:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a matrix with float32 values
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("cast_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Cast tensor to integer type
    casted_tensor = tensor.cast(DType.int32)

    print(f"Original dtype: {tensor.dtype}")  # Output: DType.float32
    print(f"Casted dtype: {casted_tensor.dtype}")  # Output: DType.int32

Parameters:: dtype (DType) – The target data type (for example, DType.int32, DType.float64).
Returns:: A new TensorValue with the casted data type.
Return type:: TensorValue

`device`

property device: DeviceRef

source

Returns the device of the TensorValue.

`dtype`

property dtype: DType

source

Returns the tensor data type.

The following example demonstrates how to access the data type of a tensor:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a matrix with float32 values
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("dtype_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Access tensor data type
    print(f"Data type: {tensor.dtype}")  # Output: DType.float32

`flatten()`

flatten(start_dim=0, end_dim=-1)

source

Flattens the specified dims of a symbolic tensor.

The number and order of the elements in the tensor is unchanged. All dimensions from start_dim to end_dim (inclusive) are merged into a single output dim.

The following example demonstrates how to flatten a multi-dimensional tensor:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x2 matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("flatten_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Flatten the tensor to a 1D array
    flattened_tensor = tensor.flatten()

    print(f"Original shape: {tensor.shape}")  # Output: [2, 2]
    print(f"Flattened shape: {flattened_tensor.shape}")  # Output: [4]

Parameters:

start_dim (int) – The starting dimension to flatten. Defaults to 0.
end_dim (int) – The ending dimension to flatten. Defaults to -1.

Returns:

A new TensorValue with the flattened dimensions.

Return type:

`from_mlir()`

classmethod from_mlir(value)

source

Creates a TensorValue from an MLIR tensor value.

Parameters:: value (Value[TensorType]) – The MLIR tensor value to wrap.
Return type:: TensorValue

`max()`

max(axis=-1)

source

Reduces the tensor using a max operation along axis.

from max.dtype import DType
from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("max_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Max along axis 1 (last dimension of each row)
    m = x.max(axis=1)

    print(f"Input shape: {x.shape}")  # [2, 3]
    print(f"Max shape: {m.shape}")    # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue with the same rank as the input and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`mean()`

mean(axis=-1)

source

Reduces the tensor using a mean operation along axis.

from max.dtype import DType
from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("mean_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Mean along axis 1 (last dimension of each row)
    mu = x.mean(axis=1)

    print(f"Input shape: {x.shape}")  # [2, 3]
    print(f"Mean shape: {mu.shape}")  # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue with the same rank as the input and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`min()`

min(axis=-1)

source

Reduces the tensor using a min operation along axis.

from max.dtype import DType

from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("min_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Min along axis 1 (last dimension of each row)
    mn = x.min(axis=1)

    print(f"Input shape: {x.shape}")  # [2, 3]
    print(f"Min shape: {mn.shape}")   # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue with the same rank as the input and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`permute()`

permute(dims)

source

Permutes the tensor’s dimensions based on provided indices.

Parameters:: dims (list[int]) – A list of integers specifying the new order of dimensions.
Returns:: A new TensorValue with permuted dimensions.
Return type:: TensorValue

`print()`

print(label='debug_tensor')

source

Prints detailed information about the tensor.

Parameters:: label (str) – A string label for the printed output. Defaults to debug_tensor.
Return type:: None

`rank`

property rank: int

source

Returns the rank (number of dims) of the buffer.

The following example demonstrates how to access the rank of a tensor:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x2 matrix (2-dimensional array)
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("rank_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Access tensor rank (number of dimensions)
    print(f"Rank: {tensor.rank}")  # Output: 2

`rebind()`

rebind(shape, message='')

source

Asserts that the tensor has the specified shape at runtime.

This method adds a runtime assertion that the tensor’s shape matches the given shape. It does not modify the tensor’s data or actual shape; instead, it verifies the shape constraint and propagates the asserted shape information to subsequent operations in the graph.

Use rebind() when you need to:

Constrain a tensor with dynamic or unknown dimensions to specific fixed sizes required by a downstream operation.
Assert shape invariants that the compiler cannot infer statically.
Provide shape hints to enable compiler optimizations.

If the assertion fails at runtime (the actual shape does not match shape), an error is raised with the optional message.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The expected shape as an iterable of integers or symbolic dimensions (Dim). The rank must match the tensor’s current rank.
message (str) – An optional message to include in the error if the assertion fails at runtime.

Returns:

A new TensorValue with the same data but with the symbolic shape set to the asserted shape.

Raises:

ValueError – If the rank of shape does not match the tensor’s rank.

Return type:

`reshape()`

reshape(shape)

source

Creates a new tensor with the same data but reshaped.

The following example demonstrates how to reshape a tensor to change its dimensions:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x2 matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("reshape_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Reshape tensor to a 1x4 matrix
    reshaped_tensor = tensor.reshape((1, 4))

    print(f"Original shape: {tensor.shape}")  # Output: [2, 2]
    print(f"Reshaped shape: {reshaped_tensor.shape}")  # Output: [1, 4]

Parameters:: shape (Iterable[int | str | Dim | integer | TypedAttr]) – The new shape as an iterable of integers or symbolic dimensions.
Returns:: A new TensorValue with the reshaped dimensions.
Return type:: TensorValue

`shape`

property shape: Shape

source

Returns the shape of the TensorValue.

The following example demonstrates how to access the shape of a tensor:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x2 matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

# Create a Graph context to work with tensors
with Graph("shape_demo") as graph:
    # Create a constant tensor from the matrix
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Access tensor shape
    print(f"Shape: {tensor.shape}")  # Shape: [Dim(2), Dim(2)]

`stdev()`

stdev(axis=-1)

source

Reduces the tensor using a standard deviation operation along axis.

The standard deviation is computed as the square root of the population variance along the specified axis.

from max.dtype import DType
from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("stdev_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Standard deviation along axis 1 (last dimension of each row)
    sd = x.stdev(axis=1)

    print(f"Input shape: {x.shape}")    # [2, 3]
    print(f"Stdev shape: {sd.shape}")  # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue with the same rank as the input and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`to()`

to(device)

source

Inserts a graph-level transfer to device into the compiled graph.

This is a graph execution-time operation: it records a transfer node during graph tracing that moves this symbolic tensor to device when the compiled graph runs. It is equivalent to calling transfer_to() and is typically used inside forward() to route activation tensors between devices.

This is distinct from to(), which is a pre-compilation host-side operation that moves stored weight tensors before the graph is built. If you want to place a module’s weights and computation on a device, use Module.to(device) before calling compile().

The following example demonstrates how to move a tensor from one device to another:

import numpy as np
from max.dtype import DType
from max.graph import Graph, ops, DeviceRef

# Create a 2x2 matrix
matrix = np.array([[1, 2], [3, 4]], dtype=np.float32)

with Graph("to_device_example") as graph:
    # Create a tensor on the default device
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Move the tensor to a GPU device
    gpu_tensor = tensor.to(DeviceRef.GPU())

    print(f"Original device: {tensor.device}")  # Output depends on default device
    print(f"New device: {gpu_tensor.device}")  # Output: gpu:0

Parameters:: device (DeviceRef) – A DeviceRef object specifying the target device.
Returns:: A new TensorValue on the specified device.
Return type:: TensorValue

`transpose()`

transpose(dim_1, dim_2)

source

Swaps two dimensions of the tensor.

The following example demonstrates how to transpose a tensor by swapping its dimensions:

import numpy as np
from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

# Create a 2x3 matrix
matrix = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)

with Graph("transpose_demo") as graph:
    tensor = ops.constant(matrix, dtype=DType.float32, device=DeviceRef.CPU())

    # Transpose the tensor (swap dimensions 0 and 1)
    transposed_tensor = tensor.transpose(dim_1=0, dim_2=1)

    print(f"Original shape: {tensor.shape}")  # Output: [2, 3]
    print(f"Transposed shape: {transposed_tensor.shape}")  # Output: [3, 2]

Parameters:

dim_1 (int) – The first dimension to swap.
dim_2 (int) – The second dimension to swap.

Returns:

A new TensorValue with swapped dimensions.

Return type:

`type`

property type: TensorType

source

Returns the type of the TensorValue as a TensorType.

`var()`

var(axis=-1)

source

Reduces the tensor using a variance operation along axis.

The variance is computed as the mean of squared deviations from the mean (population variance; that is, without Bessel’s correction) along the specified axis.

from max.dtype import DType
from max.graph import Graph, TensorType, DeviceRef

# Define a 2x3 float32 input tensor for the graph
input_type = TensorType(DType.float32, (2, 3), device=DeviceRef.CPU())
with Graph("var_demo", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor

    # Variance along axis 1 (last dimension of each row)
    vr = x.var(axis=1)

    print(f"Input shape: {x.shape}")  # [2, 3]
    print(f"Var shape: {vr.shape}")  # [2, 1]

Parameters:: axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension (for example, -1 is the last dimension).
Returns:: A TensorValue with the same rank as the input and the same shape except along axis, which will have size 1.
Return type:: TensorValue

`Type`

class max.nn.kernels.Type

source

Bases: Generic[MlirType]

The type of any value in a MAX graph.

Every value in the graph has a type, and that type is represented by a Type. This type may be inspected to get finer-grained types and learn more about an individual Value.

The following example shows how to inspect a tensor type:

from max.dtype import DType
from max.graph import DeviceRef, TensorType

tensor_type = TensorType(DType.float32, [2, 3], device=DeviceRef.CPU())
print(f"Tensor element type: {tensor_type.dtype}")  # Outputs: DType.float32
print(f"Tensor shape: {tensor_type.shape}")  # Outputs: [Dim(2), Dim(3)]

`from_mlir()`

static from_mlir(t)

source

Constructs a type from an MLIR type.

Parameters:: t (MlirType) – The MLIR Type object to parse into a type.
Returns:: The type represented by the MLIR Type value.
Return type:: Type[Any]

`to_mlir()`

to_mlir()

source

Converts to an mlir.Type instance.

Returns:: An mlir.Type in the specified Context.
Return type:: MlirType

`Value`

class max.nn.kernels.Value

source

Bases: Generic[MlirType]

Represents a symbolic value within a Graph.

A Value can represent the output of a node, the arguments of a Graph (as seen from within its body), and more generally any symbolic value available within the Graph. Other nodes receive Value values as inputs to form a computation graph.

A Value may also refer to an existing input or output of a node, and you can change them, such as by swapping a new Value.

Conceptually, think of a Value as an edge in the dataflow graph, with the other end being the user of that value.

The following example shows how to work with Values in a graph to create a simple computation:

from max.graph import DeviceRef, Graph, ops, Value
from max.dtype import DType
import numpy as np

# Create a graph context
with Graph("value_example") as graph:
    # Create input values
    a = ops.constant(np.array([1, 2, 3], dtype=np.float32), dtype=DType.float32, device=DeviceRef.CPU())
    b = ops.constant(np.array([4, 5, 6], dtype=np.float32), dtype=DType.float32, device=DeviceRef.CPU())

    # Use values to perform operations
    c = a + b  # c is a Value representing the addition

    # Demonstrate that the result is a Value
    print(f"Type of c: {type(c)}")
    print(f"Is c a Value? {isinstance(c, Value)}")

Similar to a regular variable, a Value has a data type.

Value is abstract, it shouldn’t be constructed directly.

`buffer`

property buffer: BufferValue

source

Returns the Value as a BufferValue.

Raises an exception if the Value is not a BufferValue.

`from_mlir()`

classmethod from_mlir(value)

source

Creates a Value from an MLIR value.

Parameters:: value (Value[MlirType]) – The MLIR value to wrap.
Return type:: Value[Any]

`opaque`

property opaque: _OpaqueValue

source

Returns the Value as an _OpaqueValue.

Raises an exception if the Value is not a _OpaqueValue.

`tensor`

property tensor: TensorValue

source

Returns the Value as a TensorValue.

Raises an exception if the Value is not a TensorValue.

`to_mlir()`

to_mlir()

source

Converts the Value to an MLIR value.

Return type:: Value[MlirType]

`type`

property type: Type[MlirType]

source

Returns the type of the Value as a Type.

`WeightScaleSpec`

class max.nn.kernels.WeightScaleSpec(granularity, dtype, block_size=None)

source

Bases: object

Specifies how weights are scaled for scaled quantization.

Parameters:

granularity (ScaleGranularity)
dtype (DType)
block_size (tuple[int, int] | None)

`block_size`

block_size: tuple[int, int] | None = None

source

The tuple[int, int] of the block size for block-wise scaling.

`dtype`

dtype: DType

source

The DType of the weight scale factor(s).

`granularity`

granularity: ScaleGranularity

source

The ScaleGranularity of the weight scale factor application.

`is_block`

property is_block: bool

source

Whether the weight scale granularity is block-wise.

`is_colwise`

property is_colwise: bool

source

Whether the weight scale granularity is column-wise.

`is_rowwise`

property is_rowwise: bool

source

Whether the weight scale granularity is row-wise.

`is_tensor`

property is_tensor: bool

source

Whether the weight scale granularity is per-tensor.

`accelerator_api()`

max.nn.kernels.accelerator_api()

source

Returns the API used to program the accelerator.

Return type:: str

`accelerator_architecture_name()`

max.nn.kernels.accelerator_architecture_name()

source

Returns the architecture name of the accelerator device.

Return type:: str

`amd_mega_ffn_workspace_words()`

max.nn.kernels.amd_mega_ffn_workspace_words(num_experts, max_rows, dispatch_rows)

source

Returns the u32 words of an AMD MegaFFN persistent workspace.

One padded counter slot per possible 16-row block (each expert adds at most one partial block) and per workgroup of the launch (one per CU, at most 1024), plus the two parity ticket words; then the prefill MegaFFN’s region, laid out the same way with a slot per 128-row block of the dispatch buffer (mega_prefill_region_words in mega_ffn_amd_prefill.mojo).

Parameters:

num_experts (int) – Expert slots of the FFN.
max_rows (int) – The most rows the GPU can receive inside the MegaFFN’s band.
dispatch_rows (int) – Rows of the dispatch buffer the FFN reads, its capacity at prefill.

Returns:

The workspace length in uint32 words.

Return type:

int

`apply_packed_bitmask()`

max.nn.kernels.apply_packed_bitmask(logits, packed, fill_val=-10000.0)

source

Masks logits with a packed-int32 grammar bitmask in one fused GPU pass.

Unpacks a packed bitmask (1 bit per token, 32 tokens per int32 word) and applies it to logits without materializing an intermediate bool tensor: a token is kept when its bit is set, otherwise its logit is replaced with fill_val.

Parameters:

logits (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Logits tensor of shape [batch, vocab] or [batch, num_positions, vocab].
packed (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Packed int32 bitmask of shape [..., ceil(vocab / 32)] with leading dims matching logits. A set bit means the token is grammar-valid. Trailing 32-bit alignment padding beyond vocab is never read.
fill_val (float) – Value written for masked-out (grammar-invalid) tokens.

Returns:

Masked logits, same shape and dtype as logits.

Return type:

`apply_packed_bitmask_with_penalties()`

max.nn.kernels.apply_packed_bitmask_with_penalties(logits, packed, fill_val, frequency_data, frequency_offsets, frequency_penalty, presence_penalty)

source

Applies a packed grammar bitmask, then presence/frequency penalties.

Equal to apply_packed_bitmask() followed by apply_penalties_to_logits() on the kept tokens only, in one output: a masked-out or -inf token is never penalized, so a negative penalty cannot lift it back. Zero penalties leave the mask’s output unchanged.

Parameters:

logits (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – [rows, vocab] or [batch, num_positions, vocab] logits.
packed (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Packed int32 bitmask with leading dims matching logits.
fill_val (float) – Value written for masked-out tokens.
frequency_data (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – [entries, 2] int32 CSR of [token, count]; a negative token is padding.
frequency_offsets (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – [rows + 1] uint32 start of each logit row’s entries, rows counted after flattening the leading dims.
frequency_penalty (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – [rows] float32, multiplied by the count.
presence_penalty (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – [rows] float32, subtracted once per seen token.

Returns:

Masked, penalized logits, same shape and dtype as logits.

Return type:

`apply_penalties_to_logits()`

max.nn.kernels.apply_penalties_to_logits(logits_buffer, frequency_data, frequency_offsets, *, frequency_penalty=0.0, presence_penalty=0.0, repetition_penalty=1.0)

source

Applies penalties to the logits.

Parameters:

logits_buffer (BufferValue) – The buffer to apply penalties to.
frequency_data (TensorValue) – 2d tensor of shape [unique_tokens, 2], where the first column indicates the token id and the second column indicates the frequency of the token.
frequency_offsets (TensorValue) – 1d tensor of shape [batch_size + 1], indicating start of each sequence’s data.
frequency_penalty (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The frequency penalty to apply to the model’s output. A positive value will penalize new tokens based on their frequency in the generated text: tokens will receive a penalty proportional to the count of appearances.
presence_penalty (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The presence penalty to apply to the model’s output A positive value will penalize new tokens that have already appeared in the generated text at least once by applying a constant penalty.
repetition_penalty (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The repetition penalty to apply to the model’s output. Values > 1 will penalize new tokens that have already appeared in prompt and generated text at least once by dividing the logits by the repetition penalty.

Return type:

None

`apply_qk_rms_norm()`

max.nn.kernels.apply_qk_rms_norm(q, k, qk_var, gamma_q, gamma_k, epsilon)

source

Applies QK-RMSNorm to q and k from precomputed variance.

Given the already cross-rank-reduced per-row statistics qk_var of shape [M, 2] (column 0 = mean_n(q[m, n] ** 2), column 1 = mean_n(k[m, n] ** 2), float32) and per-column float32 scales gamma_q/gamma_k, computes in a single kernel launch:

q_out[m, c] = cast((cast(q[m, c], f32) * rsqrt(qk_var[m, 0] + eps))
                   * gamma_q[c], q.dtype)
k_out[m, c] = cast((cast(k[m, c], f32) * rsqrt(qk_var[m, 1] + eps))
                   * gamma_k[c], k.dtype)

This is a single-kernel fusion of the QK-RMSNorm apply chain (rsqrt scale, gamma multiply, and downcast for both Q and K), equivalent to:

qf = ops.cast(q, DType.float32) * ops.rsqrt(qk_var[:, 0:1] + epsilon)
kf = ops.cast(k, DType.float32) * ops.rsqrt(qk_var[:, 1:2] + epsilon)
q_out = ops.cast(qf * gamma_q, q.dtype)
k_out = ops.cast(kf * gamma_k, k.dtype)

but with one launch instead of ~7 elementwise/slice kernels. The float grouping ((x * rs) * gamma) then cast matches the unfused form above. It is used for the cross-head QK-RMSNorm apply in tensor-parallel attention, where the variance is all-reduced across ranks between its computation (row_mean_of_squares_qk()) and this apply.

Parameters:

q (TensorValue) – The Q projection, [M, Nq]. bfloat16 or float32.
k (TensorValue) – The K projection, [M, Nk]. Must share q’s rows dim and dtype.
qk_var (TensorValue) – The reduced variance statistics, [M, 2] float32.
gamma_q (TensorValue) – The Q norm scale, [Nq] float32.
gamma_k (TensorValue) – The K norm scale, [Nk] float32.
epsilon (float | floating[Any]) – The RMSNorm epsilon added to the variance before rsqrt.

Returns:

A tuple (q_out, k_out) of normalized tensors matching the shapes and dtype of q and k respectively.

Raises:

ValueError – If q/k are not rank 2, have mismatched dtypes or rows, use a dtype other than bfloat16 or float32, or if the qk_var/gamma_q/gamma_k shapes or dtypes do not match.

Return type:

`argmax_with_bias()`

max.nn.kernels.argmax_with_bias(logits, bias, *, position)

source

Takes each row’s argmax of one position’s logits plus a bias.

Computes argmax(logits[:, position, :].cast(float32) + bias.cast(float32), axis=-1) in one pass, adding in float32 in registers. The upcast is exact for 16-bit inputs, so the index equals the one ops.argmax returns on the materialized float32 sum: the lowest index wins a tie and NaN never wins. logits is read in place, so a caller that walks the positions of a [rows, positions, vocab] tensor never relays it out. GPU-only.

Parameters:

logits (TensorValue) – [rows, positions, vocab_size] logits.
bias (TensorValue) – [rows, vocab_size] bias.
position (int) – The position of logits to read.

Returns:

The argmax per row, int64 [rows].

Raises:

ValueError – If the ranks are wrong, the position is out of range, or the inputs are on the CPU.

Return type:

`assert_same_device()`

max.nn.kernels.assert_same_device(*values, **named_values)

source

Raises ValueError if any of the given values are not on the same device.

Parameters:

values (TensorValue | BufferValue)
named_values (TensorValue | BufferValue)

Return type:

None

`batched_dynamic_scaled_fp8_matmul()`

max.nn.kernels.batched_dynamic_scaled_fp8_matmul(a, b, a_scales, b_scales, input_scale_spec, weight_scale_spec, out_type=bfloat16)

source

Performs a batched blockwise scaled matmul of two tensors with scaling factors.

Parameters:

a (TensorValue) – The first tensor to multiply (3D tensor).
b (TensorValue) – The second tensor to multiply, must be transposed (3D tensor).
a_scales (TensorValue) – The scaling factors for the first tensor (3D tensor).
b_scales (TensorValue) – The scaling factors for the second tensor (3D tensor).
input_scale_spec (InputScaleSpec)
weight_scale_spec (WeightScaleSpec)
out_type (DType)

Returns:

The result of the matmul operation.

Return type:

`block_diagonal_block_scaled_matmul()`

max.nn.kernels.block_diagonal_block_scaled_matmul(a, b, a_scales, b_scales, groups, tensor_sf=1.0, out_type=bfloat16)

source

Performs a block-diagonal MXFP8 block-scaled matmul on SM100.

Block g of groups multiplies a’s columns [g * k, (g + 1) * k) by b’s rows [g * n, (g + 1) * n) and writes the output’s columns [g * n, (g + 1) * n): the grouped einsum bgk,gnk->bgn with token-major operands. The blocks run as one batched kernel over strided views, so no operand is gathered or transposed.

Parameters:

a (TensorValue) – [m, groups * k] float8_e4m3fn, as quantize_dynamic_block_scaled() produces it.
b (TensorValue) – [groups * n, k] float8_e4m3fn.
a_scales (TensorValue) – a’s float8_e8m0fnu scales in the rank-5 SF-atom layout [ceildiv(m, 128), groups * k / 128, 32, 4, 4].
b_scales (TensorValue) – b’s scales in the rank-5 SF-atom layout [groups * n / 128, k / 128, 32, 4, 4].
groups (int) – The number of diagonal blocks.
tensor_sf (TensorValue | float) – Scalar multiplier applied to the product.
out_type (DType) – The output dtype.

Returns:

The [m, groups * n] product.

Raises:

ValueError – If the operands are not MXFP8 or their shapes do not split into groups blocks of whole scale tiles.

Return type:

`block_scaled_preshuffle_b_5d()`

max.nn.kernels.block_scaled_preshuffle_b_5d(b)

source

Applies the AMD CDNA4 MXFP4 B 5D preshuffle to a rank-3 weight.

Reorders the packed-FP4 bytes from [E, N, K_BYTES] row-major into the 5D (E, N0, K0, KLane=4, NLane=16, KPack=16) byte layout expected by the block_scaled_grouped_matmul_amd_preb reader. Output is byte-identical to Shuffler[E].preshuffle_b_5d running on the same input.

Intended for eager invocation from weight adapters (one-shot graph), not inside the main forward graph — the preb matmul kernel reads weights that are already in this layout.

Parameters:: b (TensorValue) – Rank-3 uint8 tensor [E, N, K_BYTES] of packed FP4 weights. N must be a multiple of 16 and K_BYTES a multiple of 64.
Returns:: Rank-3 uint8 tensor with the same shape and total byte count as b, with bytes reordered to the 5D layout.
Return type:: TensorValue

`block_scaled_preshuffle_grouped_scale_4d()`

max.nn.kernels.block_scaled_preshuffle_grouped_scale_4d(a_scales, expert_start_indices, max_num_tokens_per_expert, num_active_experts, num_experts, max_rows_per_expert=None)

source

Applies the per-step A-scale preshuffle for the AMD CDNA4 preb kernel.

Takes row-major E8M0 A-scales [total_tokens, K_SCALES] and writes cell-packed scales into per-expert fixed-stride slots of stride align_up(max_num_tokens_per_expert, 32). Output slot e holds expert slot e’s scales; the preb matmul reads from e * max_padded_M directly.

Intended to be inserted before block_scaled_grouped_matmul_amd_preb when preshuffled_b=True so the matmul sees the cell layout it expects.

Parameters:

a_scales (TensorValue) – Rank-2 float8_e8m0fnu tensor [total_tokens, K_SCALES] from quantize_dynamic_block_scaled_mxfp4. K_SCALES must be a multiple of 8.
expert_start_indices (TensorValue) – Rank-1 uint32 cumulative token offsets, length num_active_experts + 1.
max_num_tokens_per_expert (TensorValue) – Scalar uint32 upper bound on per-expert token count this step.
num_active_experts (TensorValue) – Scalar uint32 number of active expert slots.
num_experts (int) – Graph-build-time upper bound on num_active_experts (such as weight.shape[0]). Used to size the output buffer.
max_rows_per_expert (int | None) – Graph-build-time upper bound on the UNPADDED rows one expert can hold, used to size its slot. Padded here to the kernel’s 32-row stride, so pass a token count rather than a padded one. None falls back to total_tokens, which is correct but roughly top_k-fold larger per slot than anything the matmul reads. Must be positive and bound max_num_tokens_per_expert at every step: too small is silent corruption, since the matmul strides past the allocation.

Returns:

Rank-2 float8_e8m0fnu tensor [num_experts * slot_rows, K_SCALES], where slot_rows is align_up(max_rows_per_expert, 32) or total_tokens. The first num_active_experts * max_padded_M rows are written; the rest is left untouched but accessible, and is what max_rows_per_expert bounds.

Return type:

`block_scales_interleave()`

max.nn.kernels.block_scales_interleave(scales, sf_vector_size=16)

source

Interleaves rank-2 FP4 block scales into the rank-5 TCGEN layout.

Parameters:

scales (TensorValue) – Rank-2 block scales in [M, K // sf_vector_size] layout. Supported dtypes are float8_e4m3fn for NVFP4 and float8_e8m0fnu for MXFP4.
sf_vector_size (int) – Scale-factor vector size: 16 for NVFP4 or 32 for MXFP4.

Returns:

The interleaved scales tensor in [ceildiv(M, 128), ceildiv(K // sf_vector_size, 4), 32, 4, 4] layout.

Return type:

`ceildiv()`

max.nn.kernels.ceildiv(n, d)

source

Ceiling division.

Parameters:

n (Dim) – The numerator.
d (Dim) – The denominator.

Returns:

The ceiling of dividing n by d.

Return type:

Dim

`compress_windows_ragged()`

max.nn.kernels.compress_windows_ragged(kv, score, ape, norm_weight, freqs, input_row_offsets, window_offsets, starts, state, layer, *, epsilon, rotate)

source

Pools every window of a ragged chunk into one compressed entry.

DeepSeek-V4’s compressor. Request b’s chunk touches ceil(len_b / ratio) windows of ratio positions, the first at P - P % ratio for P = starts[b]. A window’s entry is a per-channel softmax over score + ape weighting kv across its positions, and its predecessor’s when ape is twice the entry’s width (the predecessor’s low half, then the window’s high half). Positions before P are read from the state leaf, later ones from kv and score. The entry is then RMS-normed, its trailing 2 * freqs.shape[1] channels rotated at the window’s first position, and snapped to the FP8 grid on the other channels, or with rotate Hadamard-rotated and snapped to FP4.

Parameters:

kv (TensorValue) – [total_rows, proj_dim] float32 raw wkv rows of the chunk.
score (TensorValue) – [total_rows, proj_dim] float32 raw wgate rows.
ape (TensorValue) – [ratio, proj_dim] float32 per-slot gate bias.
norm_weight (TensorValue) – [head_dim] RMSNorm weight, in the entries’ dtype.
freqs (TensorValue) – [positions, rope_dim // 2, 2] float32 rotary cos/sin.
input_row_offsets (TensorValue) – [batch + 1] uint32 row offsets of kv.
window_offsets (TensorValue) – [batch + 1] uint32 window offsets, request b owning ceil(len_b / ratio) entries.
starts (TensorValue) – [batch] int32 tokens cached before the chunk.
state (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The state leaf, one head of proj_dim float32 channels holding wkv as K and wgate as V by position.
layer (TensorValue) – uint32 scalar, this layer’s index in the leaf.
epsilon (float) – RMSNorm epsilon.
rotate (bool) – The indexer’s variant.

Returns:

[rows, head_dim] entries in norm_weight’s dtype, where rows = (total_rows + batch * (ratio - 1)) // ratio bounds the window count. Rows past window_offsets[batch] are left unwritten; the ragged store skips them.

Return type:

`convert_weights_to_fp8_fnuz_if_needed()`

max.nn.kernels.convert_weights_to_fp8_fnuz_if_needed(weight, weight_scale)

source

Converts weights and scales to FP8 FNUZ format if needed for AMD GPUs.

This utility function checks if FP8 FNUZ conversion is needed, currently onli AMD MI300 GPUs, and performs the conversion if required. This centralizes the conversion logic that was previously duplicated across multiple files.

Parameters:

weight (TensorValue) – The weight tensor to potentially convert.
weight_scale (TensorValue) – The weight scale factor.

Returns:

Tuple of (weight, weight_scale) - converted if needed, original otherwise.

Return type:

`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, output_dtype=None)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)
mask_variant (MHAMaskVariant)
kv_input_row_offsets (TensorValue)
q_max_seq_len (TensorValue)
scale (float)
local_window_size (int)
output_dtype (DType | None)

Return type:

`dynamic_block_scaled_matmul()`

max.nn.kernels.dynamic_block_scaled_matmul(a, b, a_scales, b_scales, tensor_sf=1.0, sf_vector_size=16, out_type=bfloat16)

source

Performs a block-scaled matmul of two NVFP4 or MXFP8 tensors.

Both formats drive the same SM100 tensor-core block-scaled MMA kernel (UMMAKind.KIND_MXF8F6F4) via the mo.matmul.dynamic.block.scaled op; the format is selected from a.dtype:

NVFP4: a/b are uint8 (two e2m1 values packed per byte) with float8_e4m3fn scales and sf_vector_size=16.
MXFP8: a/b are float8_e4m3fn (unpacked) with float8_e8m0fnu scales and sf_vector_size=32.

Parameters:

a (TensorValue) – The first tensor to multiply, rank 2 [M, K].
b (TensorValue) – The second tensor to multiply, rank 2 [N, K] (transposed).
a_scales (TensorValue) – The rank-5 SF-atom scaling factors for a.
b_scales (TensorValue) – The rank-5 SF-atom scaling factors for b.
tensor_sf (TensorValue | float) – Buffer-wise scaling factor applied to the output. For NVFP4 this is weight_scale_2 * input_scale (non-inverted); MXFP8 uses pure block scaling, so it defaults to 1.0 (identity).
sf_vector_size (int) – K-block size for the scaling factors: 16 for NVFP4 or 32 for MXFP8.
out_type (DType) – The output dtype.

Returns:

The result of the matmul operation, shape [M, N].

Return type:

`dynamic_block_scaled_matmul_amd()`

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

source

Performs a matmul of two FP4 tensors with 1D-block scaled scaling factors.

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:

`dynamic_block_scaled_matmul_mxfp6()`

max.nn.kernels.dynamic_block_scaled_matmul_mxfp6(a, b, a_scales, b_scales, fp6_format='e2m3', out_type=bfloat16, preshuffled_b=False)

source

Performs a matmul of two MXFP6 tensors with E8M0 block scales.

The FP6 sibling of dynamic_block_scaled_matmul_amd(). It is a separate op rather than another lane_bytes value because both FP6 encodings put 24 bytes in a lane, so the byte count cannot choose between them – the encoding travels as its own parameter.

AMD CDNA4 (gfx950) only.

Parameters:

a (TensorValue) – The activations, packed FP6 bytes [M, K * 3 // 4].
b (TensorValue) – The weights, packed FP6 bytes [N, K * 3 // 4] (transposed).
a_scales (TensorValue) – E8M0 activation scales [M, K // 32].
b_scales (TensorValue) – E8M0 weight scales [N, K // 32].
fp6_format (str) – The FP6 element encoding, "e2m3" or "e3m2".
out_type (DType) – The dtype of the result.
preshuffled_b (bool) – When True, b and b_scales must already be in the plane-split / packed-scale layouts produced by preshuffle_block_scaled_b() and preshuffle_block_scaled_b_scales() on a [1, N, K] view. Ignored (falls back to the row-major kernel) for small M – see mxfp6_block_scaled_matmul_amd.

Returns:

The result of the matmul operation, [M, N].

Return type:

`dynamic_scaled_matmul()`

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

source

Performs 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.
input_scale_spec (InputScaleSpec)
weight_scale_spec (WeightScaleSpec)
out_type (DType)

Returns:

The result of the matmul operation.

Return type:

`eagle_prefill_shift_tokens()`

max.nn.kernels.eagle_prefill_shift_tokens(tokens, offsets, shift_next_tokens)

source

Shifts ragged tokens left by 1 per request, appending bonus tokens.

Parameters:

tokens (TensorValue) – Flat ragged token sequence of shape [total_seq_len], dtype int64.
offsets (TensorValue) – Row offsets of shape [batch_size + 1], dtype uint32.
shift_next_tokens (TensorValue) – One token per request of shape [batch_size], dtype int64, to append after shifting.

Returns:

Shifted (or copied) tokens with the same shape as tokens.

Return type:

`flare_mla_decode_ragged()`

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

source

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

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)
mask_variant (MHAMaskVariant)
scale (float)
scalar_args (TensorValue)
qk_rope_dim (int)

Return type:

`flare_mla_decode_ragged_scaled()`

max.nn.kernels.flare_mla_decode_ragged_scaled(kv_params, input, input_row_offsets, kv_collection, kv_scales, q_scales, layer_idx, mask_variant, scale, scalar_args, qk_rope_dim=64, per_token_scale_rope_aware=False, quantization_granularity=640)

source

MLA decode with explicit per-token KV and Q scale tensors.

Like flare_mla_decode_ragged but accepts explicit scale tensors so the per-token-scale rope-aware kernel receives real (non-identity) scales.

Parameters:

kv_params (KVCacheParams) – KV cache parameters.
input (TensorValue) – Query tensor [total_tokens, num_heads, head_dim].
input_row_offsets (TensorValue) – Ragged row offsets [batch_size + 1].
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV cache collection.
kv_scales (BufferValue) – Per-token KV scales buffer [num_blocks, 1, 1, page_size, 1, 1] float32.
q_scales (TensorValue) – Per-token Q scales tensor [total_tokens] float32.
layer_idx (TensorValue) – Layer index (uint32, on CPU).
mask_variant (MHAMaskVariant) – Attention mask variant.
scale (float) – Softmax scale (typically 1/sqrt(d_qk)).
qk_rope_dim (int) – Rope head dimension (default 64).
per_token_scale_rope_aware (bool) – Use FP8+BF16 interleaved layout.
quantization_granularity (int) – Granularity for KV scale quantization. Should equal the KV cache head_dim (640 for rope-aware).
scalar_args (TensorValue)

Returns:

Output tensor [total_tokens, num_heads, output_dim].

Return type:

`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)

source

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)
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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)
buffer_size (int)

Return type:

`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)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)
buffer_size (int)
max_chunks (int)

Return type:

`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, output_dtype=None)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue]) – KV collection
layer_idx (TensorValue) – Layer index tensor
mask_variant (MHAMaskVariant) – Mask variant
scale (float) – Scale
qk_rope_dim (int) – QK rope dimension
output_dtype (DType | None) – Dtype for the attention output. Defaults to input.dtype. FP8 inputs typically pass bfloat16 since the MFMA pipeline accumulates in f32 and stores bf16.

Returns:

The output tensor for this iteration

Return type:

`flash_attention_gpu()`

max.nn.kernels.flash_attention_gpu(q, k, v, mask_variant, scale, local_window_size=-1, valid_length=None)

source

Computes flash attention using a GPU-optimized kernel.

Parameters:

q (TensorValue) – The query tensor, of shape [batch, seq_len, num_heads, head_dim].
k (TensorValue) – The key tensor, of shape [batch, seq_len, num_heads, head_dim].
v (TensorValue) – The value tensor, of shape [batch, seq_len, num_heads, head_dim].
mask_variant (MHAMaskVariant) – The mask variant to use for attention.
scale (float) – The scaling factor for attention scores.
local_window_size (int) – The local window size for sliding window attention.
valid_length (TensorValue | None) – The optional tensor of shape [batch] with dtype uint32. When provided, uses the padded kernel variant that respects the valid sequence lengths for each batch element.

Returns:

The output tensor, of shape [batch, seq_len, num_heads, head_dim].

Return type:

The following example applies causal attention to a batch of queries on a GPU:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.attention import MHAMaskVariant
from max.nn.kernels import flash_attention_gpu

batch, seq_len, num_heads, head_dim = 2, 8, 4, 128
shape = [batch, seq_len, num_heads, head_dim]
qkv_type = TensorType(DType.float32, shape, device=DeviceRef.GPU())

with Graph(
    "flash_attention_gpu_example",
    input_types=(qkv_type, qkv_type, qkv_type),
) as graph:
    # NumPy has no bfloat16, so cast at the graph boundary.
    q, k, v = (
        value.tensor.cast(DType.bfloat16) for value in graph.inputs
    )
    output = flash_attention_gpu(
        q,
        k,
        v,
        mask_variant=MHAMaskVariant.CAUSAL_MASK,
        scale=head_dim**-0.5,
    )
    graph.output(output.cast(DType.float32))

rng = np.random.default_rng(0)
q_np, k_np, v_np = (
    rng.standard_normal(shape, dtype=np.float32) for _ in range(3)
)
gpu = Accelerator()
model = InferenceSession(devices=[gpu]).load(graph)
inputs = [Buffer.from_numpy(x).to(gpu) for x in (q_np, k_np, v_np)]
result = model.execute(*inputs)[0].to_numpy()
# result.shape is (2, 8, 4, 128): one attention output per query.

`flash_attention_padded_kv_cache()`

max.nn.kernels.flash_attention_padded_kv_cache(kv_params, q, kv_collection, layer_idx, valid_lengths, mask_variant, scale, local_window_size=-1)

source

Computes flash attention with padded inputs and paged KV cache.

Parameters:

kv_params (KVCacheParams) – KV cache parameters
q (TensorValue) – Query tensor of shape [batch, seq_len, num_heads, head_dim]
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV cache collection
layer_idx (TensorValue) – Layer index for cache lookup
valid_lengths (TensorValue) – Buffer of shape [batch] with dtype uint32 indicating actual (non-padded) sequence lengths for each batch element
mask_variant (MHAMaskVariant) – The mask variant to use for attention
scale (float) – Scaling factor for attention scores
local_window_size (int) – Local window size for sliding window attention

Returns:

Output tensor of shape [batch, seq_len, num_heads, head_dim]

Raises:

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

Return type:

`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, sink_weights=None, rel_logits=None, output_dtype=None)

source

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().

When rel_logits is set, the kernel gathers an additive relative-position bias by rel_dist = q_pos - k_pos from a (total_q_tokens, heads, extent) table and adds it on every visible position, selecting mo.mha.ragged.paged.rel_logits. This path only supports mask_variant values CAUSAL_MASK (with local_window_size == -1) and SLIDING_WINDOW_CAUSAL_MASK (with a positive local_window_size); sink_weights and rel_logits are mutually exclusive.

Parameters:

kv_params (KVCacheParams) – KVCacheParams object containing key-value cache parameters.
input (TensorValue) – TensorValue representing the input tensor with shape [total_seq_len, num_heads, head_dim].
input_row_offsets (TensorValue) – TensorValue indicating the start and end of each batch in the input tensor with shape [batch_size + 1].
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – PagedCacheValues object for managing key-value cache.
layer_idx (TensorValue) – TensorValue representing the layer index, expected to have dtype uint32.
mask_variant (MHAMaskVariant) – MHAMaskVariant specifying the type of attention mask to use. With rel_logits, only CAUSAL_MASK and SLIDING_WINDOW_CAUSAL_MASK are supported.
scale (float) – float value used to scale the attention scores.
local_window_size (int) – int specifying the size of the local attention window, default is -1 for no local window.
sink_weights (TensorValue | None) – Optional tensor of shape [num_heads] containing learnable sink weights for each attention head.
rel_logits (TensorValue | None) – Optional relative-position bias table with shape [total_seq_len, num_heads, extent]; row r matches input’s own ragged-flat row convention.
output_dtype (DType | None) – Dtype for the attention output. Defaults to input.dtype.

Return type:

`flash_attention_ragged_gpu()`

max.nn.kernels.flash_attention_ragged_gpu(q, k, v, input_row_offsets, max_seq_len, mask_variant, scale, local_window_size=-1)

source

Computes flash attention for ragged inputs using a GPU-optimized kernel, without a KV cache.

Parameters:

q (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The query tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
k (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The key tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
v (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The value tensor, of shape [total_seq_len, num_heads, head_dim] (ragged).
input_row_offsets (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The buffer of shape [batch_size + 1] with dtype uint32. Indicates where each sequence starts and ends in the ragged tensors. The values should be a prefix sum (cumulative sum) of sequence lengths.
max_seq_len (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The maximum sequence length across the batch, as a rank-1 uint32 tensor on CPU.
mask_variant (MHAMaskVariant) – The mask variant to use for attention.
scale (float) – The scaling factor for attention scores.
local_window_size (int) – The local window size for sliding window attention.

Returns:

The output tensor, of shape [total_seq_len, num_heads, head_dim].

Return type:

The following example applies causal attention to a ragged batch of two sequences, of lengths 3 and 5, on a GPU:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.attention import MHAMaskVariant
from max.nn.kernels import flash_attention_ragged_gpu

seq_lens = [3, 5]
total_seq_len, num_heads, head_dim = sum(seq_lens), 4, 128
shape = [total_seq_len, num_heads, head_dim]
qkv_type = TensorType(DType.float32, shape, device=DeviceRef.GPU())

with Graph(
    "flash_attention_ragged_gpu_example",
    input_types=(
        qkv_type,
        qkv_type,
        qkv_type,
        TensorType(
            DType.uint32, [len(seq_lens) + 1], device=DeviceRef.GPU()
        ),
        TensorType(DType.uint32, [1], device=DeviceRef.CPU()),
    ),
) as graph:
    q, k, v, input_row_offsets, max_seq_len = (
        value.tensor for value in graph.inputs
    )
    # NumPy has no bfloat16, so cast at the graph boundary.
    output = flash_attention_ragged_gpu(
        q.cast(DType.bfloat16),
        k.cast(DType.bfloat16),
        v.cast(DType.bfloat16),
        input_row_offsets,
        max_seq_len,
        mask_variant=MHAMaskVariant.CAUSAL_MASK,
        scale=head_dim**-0.5,
    )
    graph.output(output.cast(DType.float32))

rng = np.random.default_rng(0)
q_np, k_np, v_np = (
    rng.standard_normal(shape, dtype=np.float32) for _ in range(3)
)
offsets = np.array([0, 3, 8], dtype=np.uint32)
max_len = np.array([max(seq_lens)], dtype=np.uint32)

gpu = Accelerator()
model = InferenceSession(devices=[gpu]).load(graph)
inputs = [
    Buffer.from_numpy(x).to(gpu) for x in (q_np, k_np, v_np, offsets)
]
# max_seq_len stays on the host.
inputs.append(Buffer.from_numpy(max_len))
result = model.execute(*inputs)[0].to_numpy()
# result.shape is (8, 4, 128): one attention output per token.

`fused_chunked_qkv_rope_store_ragged()`

max.nn.kernels.fused_chunked_qkv_rope_store_ragged(qkv, input_row_offsets, freqs_cis, kv_collection, layer_idx, value_scale, *, num_q_heads, qk_head_dim, v_head_dim, chunk_padding=0)

source

Prepares chunked QKV projections and appends to a paged KV cache.

Applies partial NeoX RoPE to Q/K, scales V in FP32 before BF16 rounding, and zero-pads each head to the cache width. Each contiguous input chunk contains its query, key, and value head regions in that order. Query rows at or past input_row_offsets[-1] are left unwritten.

Parameters:

qkv (TensorValue) – BF16 projection with shape [tokens, chunks, chunk_width].
input_row_offsets (TensorValue) – Ragged uint32 offsets, [batch_size + 1].
freqs_cis (TensorValue) – BF16 interleaved RoPE coefficients, [positions, rotary_width].
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Unquantized, non-MLA BF16 paged cache on the same GPU.
layer_idx (TensorValue) – CPU uint32 scalar selecting the cache layer.
value_scale (float) – Finite FP32 scale applied to unpadded values.
num_q_heads (int) – Total query heads across all input chunks.
qk_head_dim (int) – Unpadded query/key head width.
v_head_dim (int) – Unpadded value head width.
chunk_padding (int) – Trailing projection rows to ignore per chunk, aligned to 8.

Returns:

Padded rotated queries, [tokens, num_q_heads, cache_head_width].

Raises:

ValueError – If shapes, dtypes, devices, or cache layout are unsupported.

Return type:

`fused_dual_qk_rms_norm_rope_nvfp4_ragged()`

max.nn.kernels.fused_dual_qk_rms_norm_rope_nvfp4_ragged(main_kv_params, index_kv_params, q_main, k_main, v_main, q_index, k_index, input_row_offsets, main_kv_collection, index_kv_collection, q_main_gamma, k_main_gamma, q_index_gamma, k_index_gamma, freqs_cis, main_epsilon, index_epsilon, layer_idx, weight_offset, interleaved=True, multiply_before_cast=True, q_main_out_dtype=float8_e4m3fn, q_index_out_dtype=None)

source

Writes MiniMax-M3’s K/V and IndexK from staging into an NVFP4 main cache.

The counterpart of fused_dual_qk_rms_norm_rope_ragged for a main cache with scales, which the projection does not write. All five operands are slices of the projection output: Q, K, IndexQ and IndexK are RMS-normed and roped, V is stored as is. K and V land in the main cache quantized to NVFP4 (one E4M3 scale per 16 elements, per-tensor scale 1); IndexK lands in the index cache at its dtype, NVFP4 too when that cache is. Q and IndexQ are returned.

Parameters:

main_kv_params (KVCacheParams) – Main cache parameters; must be NVFP4.
index_kv_params (KVCacheParams) – Index-K cache parameters.
q_main (TensorValue) – Main Q [total_seq_len, n_heads, head_dim].
k_main (TensorValue) – Main K [total_seq_len, n_kv_heads, head_dim].
v_main (TensorValue) – Main V, shaped like k_main.
q_index (TensorValue) – IndexQ [total_seq_len, num_index_heads, head_dim].
k_index (TensorValue) – IndexK [total_seq_len, 1, head_dim].
input_row_offsets (TensorValue) – Ragged offsets. Dtype uint32.
main_kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The main paged cache, with scales.
index_kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The index-K paged cache.
q_main_gamma (TensorValue) – Main-Q RMSNorm weight [head_dim].
k_main_gamma (TensorValue) – Main-K RMSNorm weight [head_dim].
q_index_gamma (TensorValue) – IndexQ RMSNorm weight [head_dim].
k_index_gamma (TensorValue) – IndexK RMSNorm weight [head_dim].
freqs_cis (TensorValue) – The shared RoPE table; input dtype or float32.
main_epsilon (float | floating[Any]) – Main RMSNorm epsilon.
index_epsilon (float | floating[Any]) – Indexer RMSNorm epsilon.
layer_idx (TensorValue) – The layer index. Dtype uint32.
weight_offset (float | floating[Any]) – Constant offset added to each RMSNorm weight.
interleaved (bool) – Whether RoPE rotates adjacent pairs.
multiply_before_cast (bool) – Whether RMSNorm multiplies by the weight before rounding to the input dtype.
q_main_out_dtype (DType) – Dtype of the returned main Q.
q_index_out_dtype (DType | None) – Dtype of the returned IndexQ (default q_index.dtype).

Returns:

A tuple (q_main, q_index) of normed and roped queries.

Raises:

ValueError – On a main cache that is not NVFP4, mismatched shapes or dtypes, or gammas that do not span the head dim.

Return type:

`fused_dual_qk_rms_norm_rope_ragged()`

max.nn.kernels.fused_dual_qk_rms_norm_rope_ragged(main_kv_params, index_kv_params, main_input, index_input, input_row_offsets, main_kv_collection, index_kv_collection, q_main_gamma, k_main_gamma, q_index_gamma, k_index_gamma, freqs_cis, main_epsilon, index_epsilon, layer_idx, weight_offset, interleaved=True, multiply_before_cast=True, q_main_out_dtype=None)

source

Fuses two fused_qk_rms_norm_rope_ragged launches into one.

MiniMax-M3 sparse layers apply the fused per-head RMSNorm+RoPE op twice back to back: once for the main GQA Q / K cache and once for the lightning indexer’s IndexQ / index-K cache. Both bands read (disjoint) slices of the same combined QKV+IndexQ matmul output, share input_row_offsets, and share freqs_cis, so this runs them in a single GPU launch. Each band’s Q is returned as a separate tensor; each band’s key cache is updated in place.

All per-head RoPE geometry (dtype, freqs_cis.shape[1] rope dim, interleaved, head dim) must be identical across the two bands; a divergence trips a compile-time assert in the kernel rather than silently mis-roping a band. The two caches may differ in KV-head count, so this is bit-exact to two separate fused_qk_rms_norm_rope_ragged calls, including when q_main_out_dtype narrows the main band.

Parameters:

main_kv_params (KVCacheParams) – KV cache parameters for the main (GQA) cache.
index_kv_params (KVCacheParams) – KV cache parameters for the index-K cache.
main_input (TensorValue) – The main Q tensor [total_seq_len, n_heads, head_dim].
index_input (TensorValue) – The indexer Q tensor [total_seq_len, num_index_heads, head_dim].
input_row_offsets (TensorValue) – Ragged offsets shared by both bands. Dtype uint32.
main_kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged cache holding the main key cache.
index_kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged cache holding the index-K cache.
q_main_gamma (TensorValue) – Rank-1 main-Q RMSNorm weight (size head_dim).
k_main_gamma (TensorValue) – Rank-1 main-K RMSNorm weight (size head_dim).
q_index_gamma (TensorValue) – Rank-1 index-Q RMSNorm weight (size head_dim).
k_index_gamma (TensorValue) – Rank-1 index-K RMSNorm weight (size head_dim).
freqs_cis (TensorValue) – The shared RoPE frequency table. Its second dimension is the RoPE dim. Must share the input dtype or be float32.
main_epsilon (float | floating[Any]) – RMSNorm epsilon for the main band.
index_epsilon (float | floating[Any]) – RMSNorm epsilon for the indexer band.
layer_idx (TensorValue) – The layer index for both caches. Dtype uint32.
weight_offset (float | floating[Any]) – Constant offset added to each RMSNorm weight.
interleaved (bool) – Whether to use the interleaved RoPE pattern (both bands).
multiply_before_cast (bool) – Whether to multiply by the effective weight before rounding the RMSNorm result to input’s dtype. Governs that round only, not the q_main_out_dtype cast.
q_main_out_dtype (DType | None) – Dtype of the returned main query tensor (default main_input.dtype); the index band is not retypable. Q rounds to main_input.dtype first, so this equals a separate ops.cast.

Returns:

A tuple (q_main, q_index) of the normalized + RoPE-applied query tensors, shaped like main_input / index_input. q_index keeps index_input’s dtype; q_main is at q_main_out_dtype.

Raises:

ValueError – On invalid ranks/dtypes, mismatched gamma sizes, a gamma size that does not match its head dim, a head-dim mismatch between the two bands, or a freqs_cis dtype that neither matches the input nor is float32.

Return type:

`fused_qk_ragged_rms_norm()`

max.nn.kernels.fused_qk_ragged_rms_norm(kv_params, input, input_row_offsets, kv_collection, q_gamma, k_gamma, epsilon, layer_idx, weight_offset, multiply_before_cast=True)

source

Computes fused query-key RMSNorm with ragged inputs and paged KV cache.

This function applies per-head RMSNorm to the query tensor and to the new key entries already written into the paged KV cache. The query tensor is returned as a new tensor with the same shape and dtype as input. The key cache is normalized in place for only the newly written entries described by input_row_offsets.

Parameters:

kv_params (KVCacheParams) – The KV cache parameters.
input (TensorValue) – The query tensor of shape [total_seq_len, n_heads, head_dim].
input_row_offsets (TensorValue) – The ragged tensor offsets indicating where each batch starts and ends in input. Must have dtype uint32.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection containing the key cache.
q_gamma (TensorValue) – The rank-1 query RMSNorm weight. Its size must match kv_params.head_dim.
k_gamma (TensorValue) – The rank-1 key RMSNorm weight. Its size must match q_gamma and kv_params.head_dim.
epsilon (float | floating[Any]) – The RMSNorm epsilon value.
layer_idx (TensorValue) – The layer index for the KV cache. Must have dtype uint32.
weight_offset (float | floating[Any]) – The constant offset added to each RMSNorm weight.
multiply_before_cast (bool) – Whether to multiply by the effective weight before casting to the output dtype.

Returns:

The normalized query tensor with the same shape and dtype as input.

Raises:

ValueError – This includes when the input ranks are invalid, the row offset or layer index dtypes are invalid, the gamma weights have different sizes, or the gamma size does not match the head dimension.

Return type:

`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, position_ids=None, mrope_section=None)

source

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

Parameters:

kv_params (KVCacheParams) – KV cache parameters
input (TensorValue) – [batch_size * seq_len, n_heads, head_dim]
input_row_offsets (TensorValue) – Ragged tensor offsets indicating where each batch starts and ends
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – KV cache collection
freqs_cis (TensorValue) – tensor of shape (max_seq_len * 2, head_dim)
layer_idx (TensorValue) – Layer index for KV cache
interleaved (bool) – Whether to use interleaved RoPE pattern
position_ids (TensorValue | None) – Optional ragged 2D array of position IDs. If None, defaults to cache_length + token_idx for each token. When num_sections > 1, mrope_section must be provided to indicate each section of the head_dim to apply RoPE to. Shape: [num_sections, total_seq_len]
mrope_section (list[int] | None) – Optional list of integers indicating the section of the head_dim to
position_ids. (apply RoPE to. Must be used in conjunction with)

Return type:

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. If input is not of the same dtype as freqs_cis, it will be cast to the dtype of freqs_cis for the computation, and cast back to the original dtype after the computation is finished.

When position_ids and mrope_section are provided, it replaces the default position calculation (cache_length + token_idx) with explicit position values. This is useful for 3D RoPE in models like Qwen2.5-VL that need custom position encoding.

`fused_qk_rms_norm_rope_ragged()`

max.nn.kernels.fused_qk_rms_norm_rope_ragged(kv_params, input, input_row_offsets, kv_collection, q_gamma, k_gamma, freqs_cis, epsilon, layer_idx, weight_offset, interleaved=True, multiply_before_cast=True, q_out_dtype=None)

source

Computes fused per-head RMSNorm and RoPE with ragged inputs and paged KV cache.

This fuses fused_qk_ragged_rms_norm and fused_qk_ragged_rope into a single GPU launch. It applies per-head RMSNorm to the query tensor and to the new key entries written into the paged KV cache, then applies RoPE to the normalized values. The query tensor is returned as a new tensor with the same shape as input, at q_out_dtype (default input.dtype); the key cache is updated in place for the newly written entries.

The RoPE dimension is taken from freqs_cis.shape[1]. When it is smaller than the head dimension, RoPE is applied only to the prefix [0, rope_dim) of each head (non-interleaved layout) and the suffix is left un-roped, matching fused_qk_ragged_rope.

Parameters:

kv_params (KVCacheParams) – The KV cache parameters.
input (TensorValue) – The query tensor of shape [total_seq_len, n_heads, head_dim].
input_row_offsets (TensorValue) – The ragged tensor offsets indicating where each batch starts and ends in input. Must have dtype uint32.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection containing the key cache.
q_gamma (TensorValue) – The rank-1 query RMSNorm weight. Its size must match kv_params.head_dim.
k_gamma (TensorValue) – The rank-1 key RMSNorm weight. Its size must match q_gamma and kv_params.head_dim.
freqs_cis (TensorValue) – The RoPE frequency tensor. Its second dimension determines the RoPE dimension. Must share input’s dtype, or be float32 (the kernel upcasts freqs to the fp32 accumulator, so a higher-precision table is consumed losslessly).
epsilon (float | floating[Any]) – The RMSNorm epsilon value.
layer_idx (TensorValue) – The layer index for the KV cache. Must have dtype uint32.
weight_offset (float | floating[Any]) – The constant offset added to each RMSNorm weight.
interleaved (bool) – Whether to use the interleaved RoPE pattern.
multiply_before_cast (bool) – Whether to multiply by the effective weight before rounding the RMSNorm result to input’s dtype. Governs that round only, not the q_out_dtype cast.
q_out_dtype (DType | None) – Dtype of the returned query tensor (default input.dtype). Q rounds to input.dtype first, so this equals a separate ops.cast bit for bit.

Returns:

The normalized and RoPE-applied query tensor with the same shape as input, at q_out_dtype.

Raises:

ValueError – If the input ranks are invalid, the row offset or layer index dtypes are invalid, the gamma weights have different sizes, the gamma size does not match the head dimension, or freqs_cis dtype neither matches input nor is float32.

Return type:

`fused_qkv_padded_matmul()`

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

source

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

This is for non-ragged (padded batch) inputs where sequences may have different actual lengths but are padded to a uniform shape.

Parameters:

kv_params (KVCacheParams) – KV cache parameters.
input (TensorValue) – Input tensor with shape [batch_size, seq_len, hidden_dim].
wqkv (TensorValue) – Weight tensor for Q, K, V projections.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV cache collection.
layer_idx (TensorValue) – Layer index for cache lookup (must be uint32).
valid_lengths (TensorValue) – Buffer of shape [batch] containing the valid length for each sequence (must be uint32). K and V are only written to cache for positions within these lengths.
n_heads (int) – Number of attention heads.

Returns:

Query projections tensor. K and V projections are written to cache.

Raises:

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

Return type:

`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, _output_dim=None)

source

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

Parameters:

kv_params (KVCacheParams) – KVCacheParams object containing key-value cache parameters.
input (TensorValue) – TensorValue representing the input tensor with shape [total_seq_len, hidden_dim].
input_row_offsets (TensorValue) – TensorValue indicating the start and end of each request in the input tensor with shape [batch_size + 1].
wqkv (TensorValue) – The concatenated Q, K and V projection weights.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – PagedCacheValues object for managing key-value cache.
layer_idx (TensorValue) – TensorValue representing the layer index, expected to have dtype uint32.
n_heads (int) – Number of Query attention heads.
bias (TensorValue | None) – Optional bias vector concatenated as [q, k, v].
_output_dim (int | None) – Optional output dimension. If not provided, the output dimension will be [n_heads * head_dim].

Returns:

Query projection tensor.

Return type:

`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)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)
n_heads (int)
quantization_config (QuantizationConfig)
perm_idx (TensorValue | None)
bias (TensorValue | None)

Return type:

`grouped_dynamic_block_scaled_matmul_amd()`

max.nn.kernels.grouped_dynamic_block_scaled_matmul_amd(hidden_states, weight, a_scales, b_scales, expert_start_indices, expert_ids, expert_usage_stats_host, out_type=bfloat16, estimated_total_m=None, preshuffled_b=False, a_scales_max_padded_m=0, a_scales_max_rows_per_expert=None, decode_grid_m_cap=0, decode_grid_m_rows=0)

source

Performs grouped block-scaled AMD matmul for MoE layers.

Performs a grouped matmul with MX block-scaled inputs and weights: MXFP4 on both (packed as uint8, 2 values per byte), MXFP8 on both (float8_e4m3fn), or W4A8 with MXFP8 inputs against packed MXFP4 weights. Every format uses float8_e8m0fnu scaling factors with fixed 1D block scaling, 32 elements per scale factor along the K dimension.

hidden_states and expert_start_indices together implement the ragged tensor representation for variable-length expert inputs.

Parameters:

hidden_states (TensorValue) – The input activations with shape [total_tokens, K/2] at MXFP4 or [total_tokens, K] at MXFP8, where K is the unpacked hidden dimension. float8_e4m3fn activations are accepted against either float8_e4m3fn or packed uint8 weights.
weight (TensorValue) – The expert weights, shaped [num_experts, N, K/2] at MXFP4 or [num_experts, N, K] at MXFP8. Equal MX formats are supported; dense W4A8 uses E4M3 activations and packed uint8 MXFP4 weights.
a_scales (TensorValue) – Scaling factors for inputs with shape [num_scale_rows, K/32]. Dtype must be float8_e8m0fnu.
b_scales (TensorValue) – Scaling factors for weights with shape [num_experts, N, K/32]. Dtype must be float8_e8m0fnu.
expert_start_indices (TensorValue) – Indices indicating where each expert’s tokens start in hidden_states.
expert_ids (TensorValue) – The expert ID for each group.
expert_usage_stats_host (TensorValue) – A tensor containing [max_tokens_per_expert, num_active_experts].
out_type (DType) – Output dtype. Defaults to bfloat16.
estimated_total_m (TensorValue | None) – The estimated total number of tokens.
a_scales_max_padded_m (int) – Nonzero when a_scales already sit in the per-expert scale_4d slot layout (the EP path): the build-time slot stride their producer wrote them at. Zero means row-major A-scales, which this op preshuffles when preshuffled_b.
a_scales_max_rows_per_expert (int | None) – Graph-build-time bound on the unpadded rows one expert can hold, sizing the per-step A-scale slot buffer. None falls back to the post-expansion row count, which over-allocates each slot roughly top_k-fold. Must be positive when given, and requires preshuffled_b without a_scales_max_padded_m, the only combination that sizes slots.
decode_grid_m_cap (int) – Decode-band gate on the AMD preb path; 0 disables. Selects the band; decode_grid_m_rows bounds the grid.
decode_grid_m_rows (int) – Rows grid.y must cover per expert on the decode bands. Ignored unless preshuffled_b.
preshuffled_b (bool)

Returns:

The matmul result with shape [total_tokens, N] and dtype out_type.

Return type:

`grouped_dynamic_block_scaled_matmul_amd_swiglu_quant()`

max.nn.kernels.grouped_dynamic_block_scaled_matmul_amd_swiglu_quant(hidden_states, weight, a_scales, b_scales, expert_start_indices, expert_ids, expert_usage_stats_host, estimated_total_m=None, a_scales_max_padded_m=0, decode_grid_m_cap=0, decode_grid_m_rows=0, out_scales_max_padded_m=0, activation=None)

source

Performs a grouped MXFP8 matmul with SwiGLU and requantize fused in.

Runs gemm1 of an MoE FFN and folds the gated activation plus the MXFP8 requantize into its epilogue, so the activation leaves the kernel already quantized and no standalone activation kernel is needed. Because each (gate, up) pair collapses to one output value, the result carries half the matmul’s N.

Requires weight and b_scales sigma-permuted on N – gate and up in adjacent columns – since the epilogue pairs columns inside one lane’s accumulator fragment. AMD preshuffled-B only.

Parameters:

hidden_states (TensorValue) – MXFP8 activations, shape [total_tokens, K].
weight (TensorValue) – Sigma-permuted expert weights, [num_experts, N, K].
a_scales (TensorValue) – Activation E8M0 scales, [num_scale_rows, K/32].
b_scales (TensorValue) – Sigma-permuted weight E8M0 scales, [num_experts, N, K/32].
expert_start_indices (TensorValue) – Where each expert’s tokens start.
expert_ids (TensorValue) – The expert ID for each group.
expert_usage_stats_host (TensorValue) – [max_tokens_per_expert, num_active_experts].
estimated_total_m (TensorValue | None) – Row-count estimate used to select the band.
a_scales_max_padded_m (int) – Nonzero when the INPUT A-scales already sit in the per-expert slot layout (the KS224 ep_wait fold): the build-time slot stride the producer wrote them at, so the standalone preshuffle is skipped.
decode_grid_m_cap (int) – Decode-band gate; 0 disables.
decode_grid_m_rows (int) – Rows grid.y must cover per expert at decode.
out_scales_max_padded_m (int) – Per-expert slot stride for the EMITTED scales. Nonzero writes the 4D slot layout that gemm2 consumes without a preshuffle; zero writes plain row-major. Distinct from a_scales_max_padded_m, which describes the input.
activation (GatedActivation | None) – The gated activation applied to each (gate, up) pair: GatedActivation’s fused_activation with its alpha and limit. Defaults to plain SiLU. It must not set gated_activation_fn, which no kernel fuses.

Returns:

The MXFP8 activation [total_tokens, N/2] and its E8M0 block scales.

Return type:

`grouped_dynamic_scaled_fp8_matmul()`

max.nn.kernels.grouped_dynamic_scaled_fp8_matmul(hidden_states, weight, a_scales, b_scales, expert_start_indices, expert_ids, expert_usage_stats_host, input_scale_spec, weight_scale_spec, out_type=bfloat16, static_grid_z=False, decode_grid_m_cap=0)

source

Grouped blockwise scaled matmul used in MoE layer.

Perform a grouped blockwise scaled matmul of two tensors with scaling factors. hidden_states and expert_start_indices are used together to implement the ragged tensor.

Parameters:

hidden_states (TensorValue) – The first tensor to multiply. (2D tensor)
weight (TensorValue) – The second tensor to multiply, must be transposed. (3D tensor)
a_scales (TensorValue) – The scaling factors for the first tensor. (2D tensor)
b_scales (TensorValue) – The scaling factors for the second tensor. (3D tensor)
expert_start_indices (TensorValue) – indicates where each group starts and ends in hidden_states.
expert_ids (TensorValue) – The id of the expert for each group in hidden_states.
expert_usage_stats_host (TensorValue) – The maximum number of tokens assigned to any expert, and the number of active experts.
input_scale_spec (InputScaleSpec) – The scaling granularity for the input tensor.
weight_scale_spec (WeightScaleSpec) – The scaling granularity for the weight tensor.
static_grid_z (bool) – Make grid.z the static expert count instead of the runtime num_active_experts tensor read; capture-safe on AMD.
decode_grid_m_cap (int) – Decode-band gate on the AMD path; 0 disables. Bounds grid.y to a capture-safe upper bound derived from the call shape.
out_type (DType)

Returns:

The result of the matmul operation.

Return type:

`grouped_dynamic_scaled_mxfp6_matmul()`

max.nn.kernels.grouped_dynamic_scaled_mxfp6_matmul(hidden_states, weight, a_scales, b_scales, expert_start_indices, expert_ids, expert_usage_stats_host, fp6_format='e2m3', out_type=bfloat16, estimated_total_m=None, decode_grid_m_cap=0, decode_grid_m_rows=0, a_scales_max_padded_m=0)

source

Performs a grouped MXFP6 matmul for MoE layers.

The FP6 sibling of grouped_dynamic_scaled_mxfp4_matmul(). Both operands are packed FP6 bytes (four codes per three bytes) with E8M0 scales over 32-element K blocks.

Preshuffled-B only: an FP6 lane fragment is 24 bytes and reaches the MFMA plane-split, a layout the dense row-major grouped kernel has no path for. weight must therefore already carry the plane-split permutation from preshuffle_block_scaled_b() with lane_bytes=MXFP6_LANE_BYTES.

Parameters:

hidden_states (TensorValue) – Packed activations [total_tokens, K * 3 // 4].
weight (TensorValue) – Plane-split preshuffled expert weights [num_experts, N, K * 3 // 4].
a_scales (TensorValue) – E8M0 activation scales [num_scale_rows, K // 32].
b_scales (TensorValue) – E8M0 weight scales [num_experts, N, K // 32].
expert_start_indices (TensorValue) – Where each expert’s tokens start.
expert_ids (TensorValue) – The expert ID for each group.
expert_usage_stats_host (TensorValue) – [max_tokens_per_expert, num_active_experts].
fp6_format (str) – The FP6 element encoding, "e2m3" or "e3m2".
out_type (DType) – Output dtype.
estimated_total_m (TensorValue | None) – Estimated total token count, used to pick the persistent vs direct kernel band.
decode_grid_m_cap (int) – Decode-band gate; 0 disables.
decode_grid_m_rows (int) – Rows grid.y must cover per expert on the decode bands.
a_scales_max_padded_m (int) – Nonzero when a_scales already sit in the per-expert scale_4d slot layout: the build-time slot stride their producer wrote them at. Zero preshuffles row-major A-scales here.

Returns:

The matmul result, [total_tokens, N].

Return type:

`grouped_ffn_block_scaled_amd()`

max.nn.kernels.grouped_ffn_block_scaled_amd(hidden_states, gate_up_weight, a_scales, gate_up_scales, down_weight, down_scales, expert_start_indices, expert_ids, expert_usage_stats_host, workspace, estimated_total_m, scales_max_padded_m, decode_grid_m_cap, decode_grid_m_rows, mega_band_m, mega_max_rows, prefill_band_m=0, prefill_mega=True, activation=None)

source

Runs an MXFP8 MoE FFN on AMD gfx950: gate/up, activation, down.

In its band it is the one-launch MegaFFN; otherwise it runs the same two kernels as grouped_dynamic_block_scaled_matmul_amd_swiglu_quant() followed by grouped_dynamic_block_scaled_matmul_amd(), bit for bit either way.

Parameters:

hidden_states (TensorValue) – MXFP8 rows grouped by expert, [total_tokens, H].
gate_up_weight (TensorValue) – Preshuffled, sigma-permuted gate/up weight, [num_experts, 2 * I, H].
a_scales (TensorValue) – The rows’ E8M0 scales in the per-expert slot layout.
gate_up_scales (TensorValue) – Preshuffled gate/up weight scales.
down_weight (TensorValue) – Preshuffled down weight, [num_experts, H, I].
down_scales (TensorValue) – Preshuffled down weight scales.
expert_start_indices (TensorValue) – Where each expert’s rows start.
expert_ids (TensorValue) – The expert ID for each slot.
expert_usage_stats_host (TensorValue) – [max_tokens_per_expert, num_active_experts].
workspace (BufferValue) – The persistent MegaFFN counter buffer: minted once per call site with ops.buffer_create(..., init_value=0), sized by amd_mega_ffn_workspace_words(), never zeroed again.
estimated_total_m (TensorValue) – Row-count estimate that picks the regime.
scales_max_padded_m (int) – Per-expert slot stride of the scales.
decode_grid_m_cap (int) – Decode-band gate; 0 disables the MegaFFN.
decode_grid_m_rows (int) – Rows the chain’s decode grid covers per expert.
mega_band_m (int) – The MegaFFN runs when estimated_total_m is at most this (0: never).
mega_max_rows (int) – The most rows this GPU can receive inside that band.
prefill_band_m (int) – Above this estimated_total_m the op runs the prefill GEMMs (0: never), bit for bit the chain’s output.
prefill_mega (bool) – Runs that band as one prefill MegaFFN launch rather than two GEMM launches.
activation (GatedActivation | None) – The gated activation, as for grouped_dynamic_block_scaled_matmul_amd_swiglu_quant().

Returns:

The bf16 FFN output, [total_tokens, H].

Return type:

`grouped_matmul_block_scaled()`

max.nn.kernels.grouped_matmul_block_scaled(hidden_states, weight, a_scales, b_scales, expert_start_indices, a_scale_offsets, expert_ids, expert_scales, expert_usage_stats_host, out_type=bfloat16, estimated_total_m=None, a_row_scales=None)

source

Performs a grouped block-scaled matmul for MoE layers.

Supports four operand/scale combinations, which the op tells apart from the operand and scale dtypes alone. Every one uses fixed 1D block scaling along the K dimension:

NVFP4: both operands packed uint8, float8_e4m3fn scales over 16-element K groups.
MXFP4: both operands packed uint8, float8_e8m0fnu scales over 32-element K groups.
MXFP8: both operands float8_e4m3fn, float8_e8m0fnu scales over 32-element K groups.
W4A8: float8_e4m3fn activations against packed uint8 weights, float8_e8m0fnu scales over 32-element K groups.

Packed uint8 carries 2 4-bit E2M1 values per byte, so a packed operand’s row is K/2 wide. W4A8 is the one combination whose two operands differ, and it pairs an unpacked activation row with a packed weight row. It also requires K to be a multiple of 128, which the padded FP4 TMA copy that feeds the weights into shared memory imposes.

Every combination needs the activation row to be a multiple of 16 storage elements for its TMA descriptor, so a packed-FP4 K must be a multiple of 32. For a MoE down projection K is the expert intermediate size, and that is the constraint a tensor-parallel split of it runs into.

hidden_states and expert_start_indices together implement the ragged tensor representation for variable-length expert inputs.

Parameters:

hidden_states (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input activations with shape [total_tokens, K/2] for packed uint8 or [total_tokens, K] for float8_e4m3fn, where K is the unpacked hidden dimension.
weight (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The expert weights with shape [num_experts, N, K/2] for packed uint8 or [num_experts, N, K] for float8_e4m3fn. Sized independently of hidden_states, so W4A8 combines a [total_tokens, K] activation with a [num_experts, N, K/2] weight.
a_scales (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Scaling factors for inputs with shape [num_scale_rows, K_groups, 32, 4, 4], where K_groups is ceildiv(K, 4 * group_size) for the combination’s K group size. Dtype must be float8_e4m3fn (NVFP4) or float8_e8m0fnu (MXFP4/MXFP8/W4A8).
b_scales (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Scaling factors for weights with shape [num_experts, N_groups, K_groups, 32, 4, 4]. Dtype must match a_scales.
expert_start_indices (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Indices indicating where each expert’s tokens start in hidden_states.
a_scale_offsets (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The offsets of the input scale tiles for each expert.
expert_ids (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The expert ID for each group.
expert_scales (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Per-expert scaling factors with shape [num_experts]. Dtype must be float32. Multiplied with the matmul output in the epilogue.
expert_usage_stats_host (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – A tensor containing [max_tokens_per_expert, num_active_experts].
out_type (DType) – Output dtype. Defaults to bfloat16.
estimated_total_m (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – The estimated total number of tokens.
a_row_scales (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – Optional per-row input scales with shape [total_tokens]. Dtype must be bfloat16. Output row m is multiplied by a_row_scales[m] together with its expert scale, which undoes a per-token tensor scale applied when quantizing hidden_states. NVFP4 only.

Returns:

The matmul result with shape [total_tokens, N] and dtype out_type.

Return type:

`grouped_matmul_blocked_swiglu()`

max.nn.kernels.grouped_matmul_blocked_swiglu(hidden_states, weight, a_scales, b_scales, expert_start_indices, a_scale_offsets, expert_ids, expert_usage_stats_host, expert_scales=None, c_input_scales=None, estimated_total_m=None, activation=None, a_row_scales=None)

source

Performs fused grouped block-scaled matmul + SwiGLU for NVIDIA MoE.

Replaces the two-step chain grouped_matmul_block_scaled (BF16) -> fused_silu_quantized with a single SM100 kernel whose epilogue produces NVFP4 or MXFP8 activations and a 5D scale tile directly.

The caller must pre-permute weight and b_scales on the N axis with sigma(2i)=i, sigma(2i+1)=D+i (D = moe_dim, N = 2D) so that adjacent matmul-output columns carry (gate, up) pairs. The fused output is byte-identical to the chained reference under the kernel’s default match_bf16=True setting.

Parameters:

hidden_states (TensorValue) – The input activations. Dtype must be uint8 for packed NVFP4 or float8_e4m3fn for MXFP8.
weight (TensorValue) – The sigma-permuted expert weights with shape [num_experts, 2D, K] using the same storage dtype and packed K convention as hidden_states.
a_scales (TensorValue) – Scaling factors for inputs with shape [num_scale_rows, K_groups, 32, 4, 4]. Dtype must be float8_e4m3fn for NVFP4 or float8_e8m0fnu for MXFP8.
b_scales (TensorValue) – Scaling factors for weights with shape [num_experts, N_groups, K_groups, 32, 4, 4], with the matching sigma permutation already applied on the N axis.
expert_start_indices (TensorValue) – Per-expert token-prefix offsets (uint32, rank 1).
a_scale_offsets (TensorValue) – The offsets of the input scale tiles for each expert.
expert_ids (TensorValue) – The expert ID for each group.
expert_scales (TensorValue | None) – Per-expert scaling factors with shape [num_experts]. Dtype must be float32. Multiplied with the matmul output in the epilogue.
c_input_scales (TensorValue | None) – Per-active-expert SiLU input scale used by the NVFP4 quant epilogue. MXFP8 ignores this value.
expert_usage_stats_host (TensorValue) – A tensor containing [max_tokens_per_expert, num_active_experts].
estimated_total_m (TensorValue | None) – The estimated total number of tokens.
activation (GatedActivation | None) – The gated activation applied to each (gate, up) pair: GatedActivation’s fused_activation with its alpha and limit. Defaults to plain SiLU. It must not set gated_activation_fn, which no kernel fuses.
a_row_scales (TensorValue | None) – Optional per-row input scales with shape [total_tokens]. Dtype must be bfloat16. Row m’s gate and up values are multiplied by a_row_scales[m] together with the expert scale before the SwiGLU. NVFP4 only.

Returns:

Tuple (c_packed, c_swiglu_scales) where c_packed is packed NVFP4 uint8 with shape [total_tokens, D/2] or MXFP8 float8_e4m3fn with shape [total_tokens, D]. The scale tile’s first dim matches a_scales’s first dim since the kernel re-uses a_scale_offsets as the per-expert SF offset for the output.

Return type:

`grouped_matmul_ragged()`

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

source

Performs the grouped matmul used in the 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 is a rank-1 uint32 tensor laid out as [max_tokens_per_expert, num_active_experts] (the output of moe_create_indices).

Parameters:

hidden_states (TensorValue) – The ragged input activations.
weight (TensorValue) – The expert weights, [num_experts, N, K] (Linear convention); each group computes group @ weight.T.
expert_start_indices (TensorValue) – The start index of each group in hidden_states.
expert_ids (TensorValue) – The id of the expert for each group.
expert_usage_stats (TensorValue) – The per-expert usage stats, the output of moe_create_indices.

Returns:

The ragged matmul output.

Return type:

The following example multiplies four groups of tokens by the weights of four experts on a GPU. The groups hold 3, 1, 2, and 2 tokens and map to experts 2, 0, 3, and 1:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.kernels import grouped_matmul_ragged

num_experts, total_tokens, in_features, out_features = 4, 8, 64, 32
gpu = DeviceRef.GPU()

with Graph(
    "grouped_matmul_ragged_example",
    input_types=(
        TensorType(
            DType.float32, [total_tokens, in_features], device=gpu
        ),
        TensorType(
            DType.float32,
            [num_experts, out_features, in_features],
            device=gpu,
        ),
        TensorType(DType.uint32, [num_experts + 1], device=gpu),
        TensorType(DType.int32, [num_experts], device=gpu),
        TensorType(DType.uint32, [2], device=DeviceRef.CPU()),
    ),
) as graph:
    (
        hidden_states,
        weight,
        expert_start_indices,
        expert_ids,
        expert_usage_stats,
    ) = (value.tensor for value in graph.inputs)
    # NumPy has no bfloat16, so cast at the graph boundary.
    output = grouped_matmul_ragged(
        hidden_states.cast(DType.bfloat16),
        weight.cast(DType.bfloat16),
        expert_start_indices,
        expert_ids,
        expert_usage_stats,
    )
    graph.output(output.cast(DType.float32))

rng = np.random.default_rng(0)
hidden = rng.uniform(-1, 1, (total_tokens, in_features))
hidden = hidden.astype(np.float32)
weights = rng.uniform(-1, 1, (num_experts, out_features, in_features))
weights = weights.astype(np.float32)
starts = np.array([0, 3, 4, 6, 8], dtype=np.uint32)
ids = np.array([2, 0, 3, 1], dtype=np.int32)
# [max tokens in any group, number of groups]
usage = np.array([3, num_experts], dtype=np.uint32)

gpu = Accelerator()
model = InferenceSession(devices=[gpu]).load(graph)
inputs = [
    Buffer.from_numpy(x).to(gpu) for x in (hidden, weights, starts, ids)
]
# The usage stats stay on the host.
inputs.append(Buffer.from_numpy(usage))
result = model.execute(*inputs)[0].to_numpy()
# result.shape is (8, 32): each token's row of its expert's output.

`grouped_quantize_dynamic_block_scaled()`

max.nn.kernels.grouped_quantize_dynamic_block_scaled(input, row_offsets, scales_offsets, expert_ids, sf_tensor, sf_vector_size=16, scales_type=float8_e4m3fn, out_type=uint8, indices=None, indices_divisor=1)

source

Grouped dynamic NVFP4/MXFP4/MXFP8 quantization for MoE experts.

Quantizes a concatenated token tensor where different row ranges belong to different experts, each with its own tensor-wise scale factor.

Parameters:

input (TensorValue) – The concatenated input tensor. Shape: [total_tokens, K], dtype bfloat16.
row_offsets (TensorValue) – Cumulative token offsets per expert. Shape: [num_experts + 1], dtype uint32.
scales_offsets (TensorValue) – Per-expert scale tile offset corrections. Shape: [num_experts], dtype uint32.
expert_ids (TensorValue) – Expert ID mapping (typically identity). Shape: [num_experts], dtype int32.
sf_tensor (TensorValue) – Per-expert tensor-wise scale factors. Shape: [num_experts], dtype float32.
indices (TensorValue | None) – Optional source row indices per permuted row, fusing the MoE expert permutation gather into the quantize loads. Shape: [padded_rows], dtype int32 or uint32. When given, the input is read at row indices[i] // indices_divisor for permuted row i, and the output rows and scale tiles are sized by indices.shape[0] rather than input.shape[0].
indices_divisor (int) – The divisor applied to each entry of indices. Passing the routing slot order from moe_create_indices() with the top-k as the divisor reads each slot’s token row without a separate division kernel.
sf_vector_size (int) – The block size for the scaling factors.
scales_type (DType) – Scale factor dtype. float8_e4m3fn for NVFP4.
out_type (DType) – Output dtype. uint8 for packed FP4.

Returns:

The quantized tensor [total_tokens, K // 2] and scales in rank-5 interleaved layout [total_m_tiles, K_tiles, 32, 4, 4].

Return type:

`grouped_quantize_nvfp4_row_scaled()`

max.nn.kernels.grouped_quantize_nvfp4_row_scaled(input, row_offsets, scales_offsets, indices=None, indices_divisor=1, relu2=False)

source

Quantizes grouped rows to NVFP4, each row with its own global scale.

Fuses the row amax into the quantize kernel. The result is bit-identical to this chain of ops, which it replaces:

amax = ops.max(ops.abs(ops.cast(x, float32)), axis=-1)
row_scales = ops.cast(ops.max(amax / (6 * 448), 1e-30), bfloat16)
scaled = ops.cast(
    ops.cast(x, float32) * (1.0 / ops.cast(row_scales, float32)),
    bfloat16,
)
x_fp4, scales = grouped_quantize_dynamic_block_scaled(
    scaled, ..., sf_tensor=ones, indices=indices
)

row_scales is the a_row_scales of grouped_matmul_block_scaled(). SM100 only.

Parameters:

input (TensorValue) – The BF16 rows, [source rows, K].
row_offsets (TensorValue) – Cumulative row offsets per expert, [num_experts + 1], dtype uint32.
scales_offsets (TensorValue) – Per-expert scale tile offset corrections, [num_experts], dtype uint32.
indices (TensorValue | None) – Optional source row indices per output row, [rows], dtype int32 or uint32. Output row i quantizes input row indices[i] // indices_divisor, with its own scale. The output has indices.shape[0] rows instead of input.shape[0].
indices_divisor (int) – The divisor applied to each entry of indices.
relu2 (bool) – Whether to quantize relu(x) * relu(x), rounded to BF16, in place of x.

Returns:

The packed FP4 rows [rows, K // 2], the E4M3 block scales in the rank-5 interleaved layout [total_m_tiles, K_tiles, 32, 4, 4], and the BF16 row scales [rows].

Return type:

`gumbel_argmax_from_probs()`

max.nn.kernels.gumbel_argmax_from_probs(probs, *, seed)

source

Draws one token per row, proportionally to unnormalized probabilities.

Gumbel-max over ln(p): a zero probability can never win while the row has any positive mass, and row normalization does not matter. The Gumbel noise is generated inside the kernel from the per-row seed, so rows with equal seeds draw with equal noise – speculative decoding passes one seed per request, repeated across its draft positions, to share one noise row per request. GPU-only, and not supported on Apple GPUs, where the kernel’s block reduction cannot cover a full-sized block.

Parameters:

probs (TensorValue) – Unnormalized probabilities [rows, vocab_size], float32.
seed (TensorValue) – Per-row RNG seed [rows], uint64.

Returns:

The drawn token id per row, int64 [rows].

Raises:

ValueError – If the probabilities are not rank-2 on GPU.

Return type:

`hyper_connection_gates()`

max.nn.kernels.hyper_connection_gates(hc_proj, pre_post_comb_b, pre_post_comb_scale, hc_mult, hc_eps, hc_sinkhorn_iters)

source

Computes the Manifold-Constrained Hyper-Connections (mHC) gates.

Fuses everything an mHC site does between its stream projection and its stream collapse into one kernel: the pre and post sigmoids, the softmax that seeds the comb mixer, and the hc_sinkhorn_iters Sinkhorn-Knopp steps that project comb onto the doubly-stochastic manifold. Reference: Xie et al. 2026, section 2.2 equation 8.

hc_proj is the projection of the normalized residual streams, split along its last axis into a pre band, a post band, and a flattened hc_mult x hc_mult comb band. pre_post_comb_b is the matching concatenated bias and pre_post_comb_scale holds one scale per band. All three are float32, as are all three outputs.

Rows are independent, so a ragged batch needs no row offsets: fold the batch and sequence axes into hc_proj’s leading axis.

Parameters:

hc_proj (TensorValue) – The projected residual streams. Shape: [total_seq_len, 2 * hc_mult + hc_mult ** 2].
pre_post_comb_b (TensorValue) – Per-output bias, concatenated in pre, post, comb order. Shape: [2 * hc_mult + hc_mult ** 2].
pre_post_comb_scale (TensorValue) – Per-output scale, in pre, post, comb order. Shape: [3].
hc_mult (int) – The number of parallel residual streams. Must be a power of two whose square fits in one warp.
hc_eps (float) – Epsilon guarding the Sinkhorn divisions.
hc_sinkhorn_iters (int) – Number of Sinkhorn-Knopp iterations.

Returns:

pre: stream-collapse weights. Shape: [total_seq_len, hc_mult].
post: sublayer-output placement weights in [0, 2]. Shape: [total_seq_len, hc_mult].
comb: the row-major stream mixer. Shape: [total_seq_len, hc_mult ** 2].

Return type:

A tuple of three tensors

Raises:

ValueError – If a dtype, rank, or width does not match hc_mult, or if hc_mult / hc_sinkhorn_iters are out of range.

`indexer_score_ragged()`

max.nn.kernels.indexer_score_ragged(q, weights, input_row_offsets, base, cutoff, collection, layer, *, num_candidates)

source

Scores queries against the entries of a paged compressed leaf.

DeepSeek-V4’s lightning indexer: score[t, c] = sum_h relu(q[t, h] . k[e]) * weights[t, h] over the heads in q. The candidate axis is cap = num_candidates // 2 columns of entries closed before the chunk (column c is entry c, live while c < base[b]) followed by cap columns of the chunk’s windows (column cap + w is entry base[b] + w, live while that entry is below cutoff[t]). Dead columns are 0. Every live entry must already be in the leaf; issue the store before this op.

Parameters:

q (TensorValue) – [total_rows, num_heads, head_dim] indexer queries, ragged over the batch.
weights (TensorValue) – [total_rows, num_heads] float32 per-head weights.
input_row_offsets (TensorValue) – [batch + 1] uint32 row offsets of q.
base (TensorValue) – [batch] int32 entries each request closed before this chunk.
cutoff (TensorValue) – [total_rows] int32 entries each query may see.
collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The compressed leaf (single head, paged by entry through slots_per_page).
layer (TensorValue) – uint32 scalar, this layer’s index in the leaf.
num_candidates (int) – Width of the candidate axis, even.

Returns:

[total_rows, num_candidates] float32 scores.

Return type:

`inplace_memcpy()`

max.nn.kernels.inplace_memcpy(dst, src)

source

Copies src into dst in place.

Wraps the mo.inplace_memcpy custom op. Semantically equivalent to Buffer.inplace_copy_from, but usable from within a compiled MAX graph so the copy can be scheduled alongside other graph work.

Both operands must have the same dtype and shape. The op supports the four combinations expressible with a single DeviceContext: GPU-to-GPU on the same device, GPU-to-CPU, CPU-to-GPU, and CPU-to-CPU. Cross-GPU memcpy (different GPU ids) is rejected; use an explicit cross-device transfer for that case. The compute device is inferred from the operands: if either lives on a GPU the op is scheduled on that GPU, otherwise on CPU. :param dst: Destination buffer mutated in place. :param src: Source tensor whose contents are copied into dst.

Parameters:

dst (BufferValue)
src (TensorValue)

Return type:

None

`keyed_uniform()`

max.nn.kernels.keyed_uniform(seed)

source

Draws one uniform value in [0, 1) per row, keyed off that row’s seed.

max.graph.ops.random.uniform() reads index 0 of the graph’s seed tensor and walks the flat element index as its Philox counter, so one stream covers the whole tensor. Here each row is its own Philox key: rows with equal seeds draw equal values, and a row’s draw never depends on how many rows share the launch or on where it sits among them.

Parameters:: seed (TensorValue) – RNG seed per drawn value, uint64, of any shape.
Returns:: The drawn values, float32, shaped like seed.
Raises:: ValueError – If the seed is not a uint64 tensor.
Return type:: TensorValue

`kv_cache_gather_rows_ragged()`

max.nn.kernels.kv_cache_gather_rows_ragged(kv_collection, slots, row_offsets, layer_idx, *, key_or_value)

source

Gathers key or value rows out of a paged cache by slot.

Row r of slots addresses the request b with row_offsets[b] <= r < row_offsets[b + 1], and slots[r, j] is a slot of that request, resolved through its lookup table the way the store op resolves a token. A leaf with slots_per_page < page_size is paged by slots_per_page, so its slots are entry indices. The copy is exact, in the cache’s dtype.

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged cache leaf to read (one head per slot).
slots (TensorValue) – [rows, n] int32 slot indices. Every slot must be non-negative and inside its request’s allocated pages.
row_offsets (TensorValue) – [batch + 1] uint32 ragged offsets of the rows of slots.
layer_idx (TensorValue) – uint32 scalar, the layer to read.
key_or_value (int) – Whether to read the key or the value cache.

Returns:

[rows, n, head_dim] in the cache’s dtype.

Return type:

`kv_cache_ragged_2m_iadd()`

max.nn.kernels.kv_cache_ragged_2m_iadd(kv_params, a, kv_collection, input_row_offsets, lora_end_idx, batch_seq_len, layer_idx)

source

In-place add to paged KV cache with interleaved K/V layout.

Performs an in-place addition of new key-value projections to paged KV cache. The input tensor a uses a “2M” layout where keys and values are interleaved: rows [0, m) contain keys and rows [m, 2m) contain values, where m is the number of tokens.

Parameters:

kv_params (KVCacheParams) – KV cache configuration parameters.
a (TensorValue) – Input tensor with interleaved K/V data, shape (2*m, hidden_size) where m is the number of tokens. Rows [0, m) are keys, rows [m, 2m) are values.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection containing cache blocks, cache lengths, lookup tables, and max lengths tensors.
input_row_offsets (TensorValue) – Ragged tensor offsets indicating where each batch starts and ends
lora_end_idx (TensorValue) – End index of LoRA token portion. Marks the boundary between LoRA sequences and base model sequences in the batch.
batch_seq_len (TensorValue) – Total sequence length in the batch. Used for indexing into the value portion of a.
layer_idx (TensorValue) – The transformer layer index to update in the KV cache.

Raises:

ValueError – If a does not have rank 2.
ValueError – If input_row_offsets does not have rank 1.

Return type:

None

`kv_cache_ragged_radd()`

max.nn.kernels.kv_cache_ragged_radd(kv_params, a, kv_collection, input_row_offsets, batch_offset, layer_idx)

source

This function adds a tensor to a slice of the KVCache, sliced on the batch dimension.

This expects that the requests which should be sliced out are contiguous and in the front of the tensor, and we’re only adding to the last requests in the batch.

Parameters:

a (TensorValue) – The tensor to add to the KVCache.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The KVCache collection to add to.
input_row_offsets (TensorValue) – The offsets of the input tensor.
batch_offset (TensorValue) – The batch to start applying the r-add to.
layer_idx (int) – The layer index to add to.
kv_params (KVCacheParams)

Return type:

None

`kv_cache_store_paged_padded()`

max.nn.kernels.kv_cache_store_paged_padded(kv_collection, x_cache, valid_lengths, layer_idx, *, key_or_value)

source

Stores key or value tensor into the paged KV cache (padded inputs).

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])
x_cache (TensorValue)
valid_lengths (TensorValue)
layer_idx (TensorValue)
key_or_value (int)

Return type:

None

`kv_cache_store_paged_ragged()`

max.nn.kernels.kv_cache_store_paged_ragged(kv_collection, x_cache, input_row_offsets, layer_idx, *, key_or_value)

source

Stores key or value tensor into the paged KV cache (ragged inputs).

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection to write into.
x_cache (TensorValue) – The rank-3 tensor of new projections to store.
input_row_offsets (TensorValue) – Ragged row offsets of shape [batch + 1].
layer_idx (TensorValue) – Scalar layer index identifying which layer’s cache to write.
key_or_value (int) – Whether to store into the key or the value cache.

Return type:

None

`latent_sparse_attention_ragged()`

max.nn.kernels.latent_sparse_attention_ragged(q, input_row_offsets, comp_indices, attn_sink, swa_collection, comp_collection, layer_swa, layer_comp, *, scale, window, split_k=True, portable_p_bf16=False)

source

Sparse attention over a shared K=V latent held in two paged leaves.

DeepSeek-V4 compressed sparse attention: every head attends to the same latent row per key. A query at absolute position pos (its sequence’s cache_lengths plus its offset in the step) sees the window positions max(0, pos - window + 1) .. pos from swa_collection and the compressed entries listed for it in comp_indices from comp_collection. Both leaves must already hold the current step’s rows; issue the store ops before this one, as for every paged attention.

attn_sink enters the denominator only: the max is over the gathered scores and den += exp(sink - max).

Parameters:

q (TensorValue) – [total_rows, num_heads, head_dim] queries, ragged over the batch.
input_row_offsets (TensorValue) – [batch + 1] uint32 row offsets of q.
comp_indices (TensorValue) – [total_rows, num_comp] int32 entry indices into the compressed leaf; -1 marks an unused slot.
attn_sink (TensorValue) – [num_heads] float32 per-head sink logits.
swa_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The sliding-window leaf (single latent head, paged by token position).
comp_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The compressed leaf (single latent head, paged by entry through slots_per_page).
layer_swa (TensorValue) – uint32 scalar, this layer’s index in the window leaf.
layer_comp (TensorValue) – uint32 scalar, this layer’s index in the compressed leaf.
scale (float) – Softmax scale applied to the scores.
window (int) – Sliding window length in tokens.
split_k (bool) – Whether the SM100 decode route may split a row’s keys across CTAs. Its split partial outputs are rounded to q’s dtype before they are combined; False runs one CTA per row and rounds the output once.
portable_p_bf16 (bool) – Diagnostic. Runs the portable kernel even where the SM100 route applies, rounding each attention weight to bfloat16 before it multiplies its value row.

Returns:

[total_rows, num_heads, head_dim] in q’s dtype.

Return type:

`launch_host_func()`

max.nn.kernels.launch_host_func(payload, device)

source

Enqueues a Python callback on the device stream.

Wraps the mo.launch_host_func custom op. The callback runs on a driver thread once the stream reaches this point, after all preceding work has completed.

The payload buffer must be a CPU-resident int64[2] containing (trampoline_ptr, user_data_ptr) as returned by driver.__unsafe_pack_py_host_func.

Only supported on CUDA devices.

Parameters:

payload (BufferValue) – CPU buffer of shape [2] and dtype int64 holding the packed callback pointers.
device (DeviceRef) – GPU device on whose stream to enqueue the callback.

Return type:

None

`learnable_2d_interp_pos_emb()`

max.nn.kernels.learnable_2d_interp_pos_emb(x, weight, grid_thws, time_weight)

source

Applies learnable 2D interpolated position embedding (Kimi K2.5).

For each video described by grid_thws, bicubic-interpolates weight from (H, W) to (h, w), optionally adds temporal sincos embedding when t > 1, and adds the result element-wise to x.

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Patch embeddings of shape (L, dim).
weight (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Learnable 2D grid of shape (H, W, dim).
grid_thws (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Per-video (t, h, w) of shape (N, 3), dtype int64.
time_weight (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – 1D sincos temporal embedding of shape (num_frames, dim), dtype float32.

Returns:

Tensor of shape (L, dim) with position embeddings added.

Raises:

ValueError – On invalid input shapes or dtypes.

Return type:

`masked_flash_attention_gpu()`

max.nn.kernels.masked_flash_attention_gpu(q, k, v, mask, scale)

source

Computes flash attention using a materialized additive mask.

Parameters:

q (TensorValue) – Query tensor of shape [batch, q_seq_len, num_heads, head_dim]
k (TensorValue) – Key tensor of shape [batch, kv_seq_len, num_heads, head_dim]
v (TensorValue) – Value tensor of shape [batch, kv_seq_len, num_heads, head_dim]
mask (TensorValue) – Additive mask tensor. Rank 3 of shape [batch, q_seq_len, kv_seq_len] broadcasts across attention heads. Rank 4 of shape [batch, num_heads, q_seq_len, kv_seq_len] applies a per-head bias.
scale (float) – Scaling factor for attention scores.

Returns:

Output tensor of shape [batch, q_seq_len, num_heads, head_dim]

Return type:

`matmul_k_cache_ragged()`

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

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)

Return type:

None

`matmul_k_cache_ragged_scaled_float8()`

max.nn.kernels.matmul_k_cache_ragged_scaled_float8(kv_params, hidden_states, input_row_offsets, weight, input_scale, weight_scale, kv_collection, scales_granularity_mnk, layer_idx)

source

Computes key projections with ragged input with FP8 block scaling.

Parameters:

kv_params (KVCacheParams) – KVCacheParams object containing key-value cache parameters.
hidden_states (TensorValue) – TensorValue representing the input tensor with shape [M=total_seq_len, K=hidden_dim].
input_row_offsets (TensorValue) – TensorValue indicating the start and end of each batch in the input tensor with shape [batch_size + 1].
weight (TensorValue) – TensorValue representing the weight tensor with shape [N=num_heads, K=hidden_dim].
input_scale (TensorValue) – TensorValue representing the input scale tensor with shape [ceildiv(K / BLOCK_SIZE_K), ceildiv(M / BLOCK_SIZE_M)].
weight_scale (TensorValue) – TensorValue representing the weight scale tensor with shape [ceildiv(N / BLOCK_SIZE_N), ceildiv(K / BLOCK_SIZE_K)].
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – PagedCacheValues object for managing key-value cache.
scales_granularity_mnk (tuple[int, int, int]) – tuple[int, int, int] representing the scaling (BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K).
layer_idx (TensorValue) – TensorValue representing the layer index, expected to have dtype uint32.

Raises:

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

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)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)

Return type:

None

`matmul_static_scaled_float8()`

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

source

Performs a static-scaled float8 matrix multiplication.

Computes input @ weight.T where both tensors are float8, dequantized using the provided per-tensor CPU scalar scales before accumulation. The output is always bfloat16.

Parameters:

input (TensorValue) – Input tensor of rank 2 and dtype float8_e4m3fn or float8_e4m3fnuz.
weight (TensorValue) – Weight tensor of rank 2 and matching float8 dtype, laid out so that the K dimension matches input.shape[1].
input_scale (TensorValue) – Scalar scale factor for input (shape [] or [1]), must reside on CPU.
weight_scale (TensorValue) – Scalar scale factor for weight (shape [] or [1]), must reside on CPU.

Returns:

A TensorValue of shape [input.shape[0], weight.shape[0]] and dtype bfloat16.

Raises:

ValueError – If scale shapes are not scalar, input or weight are not rank 2, K dimensions do not match, or scales are not on CPU.

Return type:

`merge_ragged_tensors()`

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

source

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

The following example interleaves two ragged tensors that share a batch size of 2. Row a = [1, 2, 3, 4, 5, 6] with offsets [0, 2, 6] and row b = [7, 8, 9, 10] with offsets [0, 3, 4] merge into [1, 2, 7, 8, 9, 3, 4, 5, 6, 10] with offsets [0, 5, 10]:

from max.driver import Accelerator, CPU, accelerator_count
from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn.kernels import merge_ragged_tensors

device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)

with Graph(
    "merge_ragged_tensors_example",
    input_types=(
        TensorType(DType.int32, ["a_seq_len"], device=device_ref),
        TensorType(DType.uint32, ["offsets_len"], device=device_ref),
        TensorType(DType.int32, ["b_seq_len"], device=device_ref),
        TensorType(DType.uint32, ["offsets_len"], device=device_ref),
    ),
) as graph:
    a, a_row_offsets, b, b_row_offsets = (v.tensor for v in graph.inputs)
    merged_tensor, merged_row_offsets = merge_ragged_tensors(
        a, a_row_offsets, b, b_row_offsets
    )
    graph.output(merged_tensor, merged_row_offsets)

`mhc_post()`

max.nn.kernels.mhc_post(streams, y, post, comb)

source

Writes a sublayer output back into the mHC residual streams.

Computes streams'[t, k] = post[t, k] * y[t] + sum_j comb[t, j, k] * streams[t, j], i.e. post * y + comb^T @ streams per token, in float32 with one rounding to the streams’ dtype. post and comb are the float32 outputs of hyper_connection_gates(), with comb reshaped to [total_seq_len, hc_mult, hc_mult].

The kernel moves 16-byte vectors along the hidden axis and rejects inputs that are not laid out for that at execution time, as it does a token count that differs between inputs.

Parameters:

streams (TensorValue) – The incoming residual streams. Shape: [total_seq_len, hc_mult, hidden_size].
y (TensorValue) – The sublayer output, in the streams’ dtype. Shape: [total_seq_len, hidden_size].
post (TensorValue) – The weight of y in each output stream, float32. Shape: [total_seq_len, hc_mult].
comb (TensorValue) – The stream mixer, float32; comb[t, j, k] weighs input stream j in output stream k. Shape: [total_seq_len, hc_mult, hc_mult].

Returns:

The updated streams, same shape and dtype as streams.

Raises:

ValueError – If a dtype, rank or device does not match, hc_mult or hidden_size is not static or disagrees between inputs, hc_mult is not a power of two, or hidden_size is not a multiple of the kernel’s block width.

Return type:

`mhc_post_supports()`

max.nn.kernels.mhc_post_supports(hc_mult, hidden_size, dtype)

source

Returns whether mhc_post() accepts streams of this shape.

The kernel loads each token’s post and comb coefficients as one vector, so hc_mult must be a power of two, and splits the hidden axis into whole thread blocks.

Parameters:

hc_mult (int) – The number of residual streams.
hidden_size (int) – The width of one stream.
dtype (DType) – The streams’ dtype.

Returns:

True if mhc_post() accepts the shape.

Return type:

bool

`mla_decode_graph()`

max.nn.kernels.mla_decode_graph(q, kv, input_row_offsets, freqs_cis, kv_norm_gamma, w_uk, w_uv, kv_params, kv_collection, layer_idx, mask_variant, scale, epsilon, v_head_dim, scalar_args, num_partitions_scalar, *, w_uk_scale=None, w_uv_scale=None, quant_config=None, scale_granularity_override=None, sparse_indices=None, sparse_topk_lengths=None, sparse_attn_sink=None, sparse_indices_stride=None, sparse_attn_sink_inert=False, index_share=False, physical_indices=None, index_kpool=1, index_pool_topk=0, index_kpool_select_tail=True, selection_lag=0)

source

This is a manually fused kernel that performs the following operations:

Apply RoPE to the query and the key cache (in-place).
Apply RMSNorm to the non-rope portion of the key cache (in-place).
Project q_nope to kv_latent_dim through a fp8 batched matmul: q_nope_proj = q_nope_t @ w_uk
Concatenate q_nope_proj and q_rope: q_full = concat(q_nope_proj, q_rope, axis=2)
Perform MLA decode
Project raw_output to v_head_dim through another fp8 batched matmul: output = raw_output_t @ w_uv

Parameters:

q (TensorValue) – Combined query tensor containing both nope and rope parts. Shape: [tot_seq_len, num_heads, qk_nope_head_dim + qk_rope_head_dim].
kv (TensorValue) – KV latent tensor from the first projection. Shape: [num_tokens, cache_head_dim] where cache_head_dim = kv_lora_rank + qk_rope_head_dim.
input_row_offsets (TensorValue) – Indicates where each request starts and ends in input. This is a 1D tensor of shape [num_batches + 1].
freqs_cis (TensorValue) – Precomputed RoPE frequency values for rotary position embeddings. Shape: [max_seq_len, qk_rope_head_dim].
kv_a_proj_layernorm – RMSNorm gamma weights for normalizing the KV cache. Shape: [kv_lora_rank].
w_uk (TensorValue) – Weight matrix for projecting q_nope to kv_latent_dim. Shape: [num_heads, kv_latent_dim, qk_nope_head_dim].
w_uv (TensorValue) – Weight matrix for projecting MLA decode output to v_head_dim. Shape: [num_heads, v_head_dim, kv_latent_dim].
kv_params (KVCacheParams) – KVCacheParams
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV Cache object.
layer_idx (TensorValue) – Layer index.
mask_variant (MHAMaskVariant) – The attention mask variant controlling masking behavior.
scale (float) – Scale for the attention calculation.
epsilon (float) – Small constant for numerical stability in RMSNorm.
v_head_dim (int) – Dimension of the V heads.
scalar_args (TensorValue) – Pre-computed dispatch scalar args (GPU buffer) for CUDA graph capture.
w_uk_scale (TensorValue | None) – Optional FP8 scale tensor for w_uk.
w_uv_scale (TensorValue | None) – Optional FP8 scale tensor for w_uv.
quant_config (QuantConfig | None) – Optional quantization config. When set, scales are required.
sparse_indices (TensorValue | None) – Optional int32 tensor of shape [total_seq_len, max_topk] selecting each token’s keys, -1 for an unused slot. physical_indices says which kind they are.
sparse_topk_lengths (TensorValue | None) – Per-batch valid top-k counts, int32 rank-1. None means every row uses the full sparse_indices_stride.
sparse_attn_sink (TensorValue | None) – Per-batch attention sink weights, float32 rank-1. None means no sink.
sparse_indices_stride (int | None) – Row stride in sparse_indices (max top-k across the batch). Required when sparse_indices is set. Over an NVFP4 cache, whose scales are paged apart from its values, a row is the top-k physical rows followed by as many scale-pool rows, as mla_fp8_index_top_k_physical() returns them, so the stride is twice the top-k.
index_share (bool) – Whether the folded MTP query positions share one top-k list, so the FP8 sparse decode gathers it once. The fold masks each position by logical key position, so it requires logical sparse_indices.
physical_indices (bool | None) – Whether sparse_indices holds physical MLA cache rows, block_id * page_size + offset, as mla_fp8_index_top_k_physical() returns them, rather than logical token positions, which the kernel looks up in the cache’s page table. Required with sparse_indices.
index_kpool (int) – Keys per pool in the k-pool indexer that produced sparse_indices; 1 for a plain top-k list. With index_pool_topk, the kernel derives each row’s valid length from the token’s visible key count and masks the -1 padding past it without reading the indices.
index_pool_topk (int) – Maximum number of pools the indexer selects; 0 for no bound beyond the row width.
index_kpool_select_tail (bool) – Whether the indexer appends the incomplete trailing pool to each row.
selection_lag (int) – Positions between the token each row of sparse_indices was computed for and the token that reads it, for a selection reused at a later position.
kv_norm_gamma (TensorValue)
num_partitions_scalar (TensorValue)
scale_granularity_override (int | None)
sparse_attn_sink_inert (bool)

Returns:

Tensor of shape [total_seq_len, num_heads, v_head_dim].

Return type:

`mla_fp8_index_top_k()`

max.nn.kernels.mla_fp8_index_top_k(q, q_s, input_row_offsets, k_collection, layer_idx, top_k, quantization_granularity, mask_variant=MHAMaskVariant.CAUSAL_MASK, kpool=1)

source

Computes top-k indices for MLA FP8 indexed attention scores.

This function computes FP8 matmul between queries and cached keys (with scales), applies masking, and returns the indices of the top-k highest-scoring keys per token. Scores are aggregated (summed) across all attention heads.

Parameters:

q (TensorValue) – Query tensor of shape [total_seq_len, num_heads, head_dim] in FP8.
q_s (TensorValue) – Query scales tensor of shape [total_seq_len, num_heads] in float32.
input_row_offsets (TensorValue) – Input row offsets tensor of shape [batch_size + 1].
k_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV cache collection. Must be FP8 quantized with scales.
layer_idx (TensorValue) – Layer index for cache lookup.
top_k (int) – Requested number of top indices per token, at most 2048.
quantization_granularity (int) – Quantization granularity for the K cache.
mask_variant (MHAMaskVariant) – The mask variant to use (NULL or CAUSAL_MASK).
kpool (int) – Tokens per row of k_collection. 1 scores one row per cached token. k > 1 means the cache holds one pooled key per k consecutive tokens, so top_k counts pools and the returned indices are pool ids that the caller expands back to token positions.

Returns:

Output tensor of shape [total_seq_len, top_k] containing top-k key indices per token. Positions past a token’s candidate keys are filled with -1. With kpool > 1 the entries are pool ids rather than token positions.

Raises:

ValueError – If top_k is not in [1, 2048], kpool is not positive, or mask_variant is neither NULL nor CAUSAL.

Return type:

`mla_fp8_index_top_k_physical()`

max.nn.kernels.mla_fp8_index_top_k_physical(q, q_s, input_row_offsets, k_collection, mla_collection, layer_idx, top_k, quantization_granularity, mask_variant=MHAMaskVariant.CAUSAL_MASK)

source

Computes only the physical rows of mla_fp8_index_top_k_with_physical().

The top-k skips writing the logical indices, which the sparse MLA ops do not read once they have the physical rows.

Parameters:

q (TensorValue) – Query tensor of shape [total_seq_len, num_heads, head_dim] in FP8.
q_s (TensorValue) – Query scales tensor of shape [total_seq_len, num_heads] in float32.
input_row_offsets (TensorValue) – Input row offsets tensor of shape [batch_size + 1].
k_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The indexer’s paged K cache. Must be FP8 with scales.
mla_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The MLA cache the physical rows index.
layer_idx (TensorValue) – Layer index for cache lookup.
top_k (int) – Requested number of top indices per token, at most 2048.
quantization_granularity (int) – Quantization granularity for the K cache.
mask_variant (MHAMaskVariant) – The mask variant to use (NULL or CAUSAL_MASK).

Returns:

A [total_seq_len, top_k] tensor of the selected positions’ physical rows in mla_collection, -1 past a token’s candidates or where a page is unmapped. For an NVFP4 mla_collection the tensor is [total_seq_len, 2 * top_k]: the physical rows, then each selection’s row in the scale pool.

Raises:

ValueError – If top_k is not in [1, 2048] or mask_variant is neither NULL nor CAUSAL.

Return type:

`mla_fp8_index_top_k_with_physical()`

max.nn.kernels.mla_fp8_index_top_k_with_physical(q, q_s, input_row_offsets, k_collection, mla_collection, layer_idx, top_k, quantization_granularity, mask_variant=MHAMaskVariant.CAUSAL_MASK)

source

Computes mla_fp8_index_top_k() and the selections’ MLA cache rows.

The second output is each selected position’s physical row in mla_collection, block_id * page_size + offset, written by the indexer’s top-k as it selects. Pass it to the sparse MLA ops as sparse_indices with physical_indices=True so they skip remapping the logical indices.

Parameters:

q (TensorValue) – Query tensor of shape [total_seq_len, num_heads, head_dim] in FP8.
q_s (TensorValue) – Query scales tensor of shape [total_seq_len, num_heads] in float32.
input_row_offsets (TensorValue) – Input row offsets tensor of shape [batch_size + 1].
k_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The indexer’s paged K cache. Must be FP8 with scales.
mla_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The MLA cache the physical rows index.
layer_idx (TensorValue) – Layer index for cache lookup.
top_k (int) – Requested number of top indices per token, at most 2048.
quantization_granularity (int) – Quantization granularity for the K cache.
mask_variant (MHAMaskVariant) – The mask variant to use (NULL or CAUSAL_MASK).

Returns:

The logical indices, as mla_fp8_index_top_k() returns them, and a tensor of the same shape holding their physical rows, -1 where the logical index is -1 or its page is unmapped. For an NVFP4 mla_collection the rows are twice as wide: the physical rows, then each selection’s row in the scale pool, which its scales lookup table pages apart from the values.

Raises:

ValueError – If top_k is not in [1, 2048] or mask_variant is neither NULL nor CAUSAL.

Return type:

`mla_kpool_compress()`

max.nn.kernels.mla_kpool_compress(k, gate, ape, input_row_offsets, pool_row_offsets, cache_lengths, kpool)

source

Compresses each complete k-pool into one candidate key.

A pool covers kpool consecutive tokens of one request and collapses to a single key: a per-channel softmax over gate + ape across the pool’s members, weighting their keys. Pooling is what lets the DSA indexer score kv_len / kpool candidates instead of kv_len; a pooled key is immutable once its members exist, which is what makes it cacheable.

Only complete pools are emitted. The tokens after the last complete pool are the caller’s tail to handle.

Parameters:

k (TensorValue) – Layer-normed indexer keys, [total_tokens, head_dim].
gate (TensorValue) – Per-token gate scores, [total_tokens, head_dim], same dtype as k.
ape (TensorValue) – Within-pool position embedding, [kpool, head_dim], float32.
input_row_offsets (TensorValue) – Token row offsets per request, [batch + 1].
pool_row_offsets (TensorValue) – Pool row offsets per request, [batch + 1].
cache_lengths (TensorValue) – Cached-prefix length per request, [batch]. A pool covers absolute positions, so this is what places the call’s tokens on the pool grid.
kpool (int) – Tokens per pool.

Returns:

[total_tokens // kpool, head_dim] pooled keys, dtype of k.

The row count is an upper bound, not the exact pool count: the exact count is data-dependent – a sum of per-request floors – and the graph compiler needs an output shape derivable from the inputs. The kernel grids on the true pool count it reads from pool_row_offsets, so trailing rows are never written and never read; pool_row_offsets is what locates each request’s pools.

Raises:

ValueError – If kpool is not positive, if the operand ranks or head dims disagree, or if ape is not [kpool, head_dim].

Return type:

`mla_kpool_expand_topk()`

max.nn.kernels.mla_kpool_expand_topk(pool_ids, input_row_offsets, cache_lengths, kpool, always_select_tail=True)

source

Expands selected k-pool ids back into the token positions they cover.

The indexer selects pools; attention reads tokens. Pool p covers positions [p * kpool, (p + 1) * kpool), so the selection widens from pool_topk to pool_topk * kpool.

With always_select_tail the row also carries up to kpool - 1 of the query’s most recent positions, the ones no complete pool covers yet, so the result is kpool - 1 columns wider. The tail is located from the query’s visible count, so it tracks the pool currently filling.

Each row is packed: the selected pools in selection order, then the tail, then -1 in every remaining column, never a clamped valid position, which would point attention at a token the indexer did not choose. Packing keeps every valid entry ahead of the padding, which the sparse decode relies on when it bounds its scan by the context length.

Parameters:

pool_ids (TensorValue) – Selected pool ids, [total_seq_len, pool_topk] int32.
input_row_offsets (TensorValue) – Token row offsets per request, [batch + 1].
cache_lengths (TensorValue) – Cached tokens per request, [batch].
kpool (int) – Tokens per pool.
always_select_tail (bool) – Whether to append the trailing partial pool.

Returns:

[total_seq_len, pool_topk * kpool + tail] int32 token positions, -1 where unused, where tail is kpool - 1 when always_select_tail and 0 otherwise.

Raises:

ValueError – If kpool is not positive or pool_ids is not rank 2.

Return type:

`mla_kpool_ring_close()`

max.nn.kernels.mla_kpool_ring_close(tail, k, gate, ape, input_row_offsets, cache_lengths, slot_idx, kpool)

source

Closes each request’s pending tail-ring pool, ragged and unconditional.

A cached prefix ending mid-pool leaves that pool’s earlier members in tail; neither mla_kpool_compress() nor mla_kpool_seed_tail() reads them back – compress only ever builds pools entirely from its own call’s tokens, skipping the leading remainder the ring already holds. This op is the one that reads that remainder, closing the pool once a call brings enough new tokens to finish it.

This takes a ragged, arbitrary-width chunk per request rather than a fixed next_n, and closes at most one pool per request per call: the ring never holds more than kpool - 1 members between calls, so there is never a second one pending to close. Any further pools a call completes are entirely new tokens and are mla_kpool_compress()’s job, which already skips exactly the tokens this op consumes.

Parameters:

tail (BufferValue) – Persistent slot-indexed ring, dtype of k, [max_slots, 2, kpool, head_dim]. Read only here – mla_kpool_seed_tail() is what mutates it.
k (TensorValue) – This call’s layer-normed keys, [total_tokens, head_dim].
gate (TensorValue) – This call’s gate scores, [total_tokens, head_dim], same dtype as k.
ape (TensorValue) – Within-pool position embedding, [kpool, head_dim], float32.
input_row_offsets (TensorValue) – Token row offsets per request, [batch + 1].
cache_lengths (TensorValue) – Cached-prefix length per request, [batch].
slot_idx (TensorValue) – Ring slot per request, [batch] uint32.
kpool (int) – Tokens per pool.

Returns:

pooled: [batch, head_dim] pooled keys, dtype of k, meaningful only where closed_pool is non-negative.
closed_pool: [batch] int32 pool id closed this call, or -1 where that request had nothing pending, or not yet enough new tokens to finish it.

Return type:

A (pooled, closed_pool) pair

Raises:

ValueError – If kpool is not positive or the operand ranks or head dims disagree.

`mla_kpool_seed_tail()`

max.nn.kernels.mla_kpool_seed_tail(tail, k, gate, input_row_offsets, cache_lengths, slot_idx, kpool)

source

Stashes a prefill chunk’s trailing tokens into the tail ring.

Whole pools compress directly from k/gate via mla_kpool_compress(); the tokens after the last complete pool are that request’s in-progress pool and have nowhere else to go until later tokens complete it, so they seed this per-request ring. mla_kpool_ring_close() reads them back and closes the pool once it fills.

Parameters:

tail (BufferValue) – Persistent slot-indexed ring, mutated in place, dtype of k, [max_slots, 2, kpool, head_dim]. Index 0 holds keys, index 1 holds gate scores.
k (TensorValue) – Layer-normed indexer keys, [total_tokens, head_dim].
gate (TensorValue) – Per-token gate scores, [total_tokens, head_dim], same dtype as k.
input_row_offsets (TensorValue) – Token row offsets per request, [batch + 1].
cache_lengths (TensorValue) – Cached-prefix length per request, [batch].
slot_idx (TensorValue) – Ring slot per request, [batch] uint32.
kpool (int) – Tokens per pool.

Raises:

ValueError – If kpool is not positive or the operand ranks or head dims disagree.

Return type:

None

`mla_prefill_decode_graph()`

max.nn.kernels.mla_prefill_decode_graph(q, kv, input_row_offsets, freqs_cis, kv_norm_gamma, buffer_row_offsets, cache_offsets, buffer_length, w_k, w_uk, w_uv, kv_params, kv_collection, layer_idx, mask_variant, scale, epsilon, v_head_dim, scalar_args, num_partitions_scalar, *, w_k_scale=None, w_uk_scale=None, w_uv_scale=None, quant_config=None, scale_granularity_override=None, sparse_indices=None, sparse_topk_lengths=None, sparse_attn_sink=None, sparse_indices_stride=None, physical_indices=None, index_kpool=1, index_pool_topk=0, index_kpool_select_tail=True, selection_lag=0, sparse_attn_sink_inert=False)

source

Fused MLA prefill/decode kernel for FP8.

Switches between prefill and decode based on the maximum sequence length in the batch. See mla_prefill_graph and mla_decode_graph for the dedicated paths.

Parameters:

q (TensorValue) – Combined query tensor with nope+rope parts.
kv (TensorValue) – KV latent tensor for current sequence.
input_row_offsets (TensorValue) – Row offsets for the batch.
freqs_cis (TensorValue) – RoPE frequencies tensor.
kv_norm_gamma (TensorValue) – RMSNorm gamma for KV cache.
buffer_row_offsets (TensorValue) – One-shot prefill buffer row offsets.
cache_offsets (TensorValue) – One-shot prefill cache offsets.
buffer_length (TensorValue) – One-shot prefill buffer length tensor.
w_k (TensorValue) – Prefill K up-projection weights.
w_uk (TensorValue) – Decode query-projection weights.
w_uv (TensorValue) – Decode output-projection / prefill V-projection weights.
kv_params (KVCacheParams) – KV cache parameters.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV cache values.
layer_idx (TensorValue) – Layer index (uint32).
mask_variant (MHAMaskVariant) – Attention mask variant.
scale (float) – Attention scale.
epsilon (float) – RMSNorm epsilon.
v_head_dim (int) – Value head dimension for output tensor shape.
scalar_args (TensorValue) – Pre-computed dispatch scalar args (GPU buffer) for CUDA graph capture.
w_k_scale (TensorValue | None) – Optional FP8 scale tensor for w_k.
w_uk_scale (TensorValue | None) – Optional FP8 scale tensor for w_uk.
w_uv_scale (TensorValue | None) – Optional FP8 scale tensor for w_uv.
quant_config (QuantConfig | None) – Optional quantization config. When set, scales are required.
sparse_indices (TensorValue | None) – Optional int32 tensor selecting each token’s keys (same semantics as mla_decode_graph()).
sparse_topk_lengths (TensorValue | None) – Per-batch valid top-k counts for sparse decode. None means every row uses the full sparse_indices_stride.
sparse_attn_sink (TensorValue | None) – Per-batch attention sink weights for sparse decode. None means no sink.
sparse_indices_stride (int | None) – Row stride in sparse_indices. Required when sparse_indices is set.
physical_indices (bool | None) – Whether sparse_indices holds physical MLA cache rows rather than logical token positions, which a decode call looks up in the page table and a prefill call remaps first. Required with sparse_indices.
index_kpool (int) – Keys per pool in the indexer that produced sparse_indices (same semantics as mla_decode_graph()).
index_pool_topk (int) – Maximum number of pools the indexer selects.
index_kpool_select_tail (bool) – Whether the indexer appends the incomplete trailing pool to each row.
selection_lag (int) – Positions between the token each row of sparse_indices was computed for and the token that reads it. Only the decode branch reads it.
num_partitions_scalar (TensorValue)
scale_granularity_override (int | None)
sparse_attn_sink_inert (bool)

Returns:

Tensor of shape [total_seq_len, num_heads, v_head_dim].

Return type:

`mla_prefill_graph()`

max.nn.kernels.mla_prefill_graph(q, kv, input_row_offsets, freqs_cis, kv_norm_gamma, buffer_row_offsets, cache_offsets, buffer_length, w_k, w_uv, kv_params, kv_collection, layer_idx, mask_variant, scale, epsilon, v_head_dim, *, w_k_scale=None, w_uv_scale=None, quant_config=None, scale_granularity_override=None)

source

This is a manually fused kernel that performs the following operations:

Apply RoPE to the query and the key cache (in-place).
Apply RMSNorm to the non-rope portion of the key cache (in-place).
Copy the KV latent values from PagedKVCache to a contiguous buffer.
Quantize the KV latent values to fp8.
Up-project the latent KV values to full K and V through two matmuls.
Perform MLA prefill.

Parameters:

q (TensorValue) – Combined query tensor containing both nope and rope parts. Shape: [tot_seq_len, num_heads, qk_nope_head_dim + qk_rope_head_dim].
kv (TensorValue) – KV latent tensor from the first projection. Shape: [num_tokens, cache_head_dim] where cache_head_dim = kv_lora_rank + qk_rope_head_dim.
input_row_offsets (TensorValue) – Indicates where each request starts and ends in input. This is a 1D tensor of shape [num_batches + 1].
freqs_cis (TensorValue) – Precomputed RoPE frequency values for rotary position embeddings. Shape: [max_seq_len, qk_rope_head_dim].
kv_a_proj_layernorm – RMSNorm gamma weights for normalizing the KV cache. Shape: [kv_lora_rank].
buffer_row_offsets (TensorValue) – Indicates where each request’s KV latent values should be stored in the contiguous buffer. This is a 1D tensor of shape [num_batches + 1].
cache_offsets (TensorValue) – Indicates the starting token position in the KV cache from which to copy KV latent values for each request. This is a 1D tensor of shape [num_batches + 1].
buffer_length (TensorValue) – The total number of tokens in the KV cache. Scalar.
w_k (TensorValue) – Weight matrix for up-projecting latent KV values to full K. Shape: [num_heads * qk_nope_head_dim, kv_latent_dim].
w_uv (TensorValue) – Weight tensor for up-projecting latent KV values to full V. Shape: [num_heads, v_head_dim, kv_latent_dim].
kv_params (KVCacheParams) – KVCacheParams
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged KV Cache object.
layer_idx (TensorValue) – Layer index.
mask_variant (MHAMaskVariant) – The attention mask variant controlling masking behavior.
scale (float) – Scale for the attention calculation.
epsilon (float) – Small constant for numerical stability in RMSNorm.
v_head_dim (int) – Dimension of the V heads.
w_k_scale (TensorValue | None) – Optional FP8 scale tensor for w_k.
w_uv_scale (TensorValue | None) – Optional FP8 scale tensor for w_uv.
quant_config (QuantConfig | None) – Optional quantization config. When set, scales are required.
kv_norm_gamma (TensorValue)
scale_granularity_override (int | None)

Returns:

Tensor of shape [total_seq_len, num_heads, v_head_dim].

Return type:

`moe_create_indices()`

max.nn.kernels.moe_create_indices(topk_ids, num_local_experts, *, needs_scales_offset=False)

source

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.
needs_scales_offset (bool)

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 that restore the original token ordering after expert computation.
expert_ids: The expert ID of each group in the reordered sequence.
expert_usage_stats: The maximum number of tokens assigned to any expert, and the number of expert groups (num_local_experts), including empty ones.

Return type:

A tuple of five tensors

The following example groups eight token assignments by expert, for eight experts on a GPU:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.kernels import moe_create_indices

num_experts = 8
topk_ids = np.array([3, 0, 3, 5, 0, 3, 7, 5], dtype=np.int32)

with Graph(
    "moe_create_indices_example",
    input_types=(
        TensorType(
            DType.int32, [len(topk_ids)], device=DeviceRef.GPU()
        ),
    ),
) as graph:
    outputs = moe_create_indices(
        graph.inputs[0].tensor, num_local_experts=num_experts
    )
    graph.output(*outputs)

gpu = Accelerator()
model = InferenceSession(devices=[gpu]).load(graph)
outputs = model.execute(Buffer.from_numpy(topk_ids).to(gpu))
(
    token_expert_order,
    expert_start_indices,
    restore_token_order,
    expert_ids,
    expert_usage_stats,
) = (buffer.to_numpy() for buffer in outputs)
# expert_start_indices is [0, 2, 2, 2, 5, 5, 7, 7, 8]: expert 0
# owns positions 0-1, expert 3 owns 2-4, expert 5 owns 5-6, and
# expert 7 owns position 7. expert_usage_stats is [3, 8].

`moe_create_indices_with_slots()`

max.nn.kernels.moe_create_indices_with_slots(topk_ids, slot_weights, num_groups, *, slot_stride, scale_rows, expert_offset=0, local_experts=None)

source

moe_create_indices() plus a slot-packed grouped GEMM’s row maps.

For a grouped GEMM over the routed slots sorted by group and packed, that reads its A-side scales from 128-row tiles per group (the layout of needs_scales_offset). The maps it would otherwise derive from the CSR in the graph come out of the same kernel, with the same values.

Expert ids are mapped to groups first: id - expert_offset when that is in [0, local_experts), else local_experts. With num_groups == local_experts + 1 that last group collects the slots of experts held on other devices (expert parallel), so every slot still gets order and restore entries.

Parameters:

topk_ids (TensorValue) – [slots] int32 expert id of each routed slot.
slot_weights (TensorValue) – [slots] float32 routing weight of each slot.
num_groups (int) – Groups to bucket into, local_experts or one more.
slot_stride (int) – Slots per token (n_experts_per_tok).
scale_rows (int | str | Dim | integer | TypedAttr) – Length of scale_slot, at least the aligned total of the group tiles; slots + 128 * num_groups always is, and so is 128 * (ceil(slots / 128) + min(num_groups, slots)).
expert_offset (int) – Global id of the first local expert.
local_experts (int | None) – Experts held locally; num_groups if None.

Returns:

token_expert_order, expert_start_indices, restore_token_order (int32), expert_ids and expert_usage_stats as moe_create_indices(), then:

scales_offset: [num_groups] uint32, as needs_scales_offset.
slot_token: [slots] int32, token_expert_order // slot_stride.
slot_weight: [slots] float32, slot_weights[token_expert_order].
scale_slot: [scale_rows] int32, the sorted slot whose scales each scale row holds: row r of group g’s tiles holds slot expert_start_indices[g] + r while r is below the group’s count, and every other row holds slot 0.

Return type:

tuple[TensorValue, …]

`moe_eplb_remap()`

max.nn.kernels.moe_eplb_remap(router_idx, logcnt, log2phy, layer_idx, *, num_log, max_replicas, n_experts_per_tok, hash_decorrelate=False)

source

Fused EPLB logical-to-physical id remap. single Mojo kernel that caches the per-layer slice of logcnt and log2phy in shared memory and writes physical ids in one launch.

The replica picker is deterministic position-mod (r = (n*K + k) % logcnt[layer, log]), bit-identical to the legacy chain when hash_decorrelate=False. With hash_decorrelate=True the flat position is xor-hashed with a Knuth multiplicative hash of the logical id before the modulo, breaking structured position-vs-cnt alignment without warp primitives.

Parameters:

router_idx (TensorValue) – [num_tokens, n_experts_per_tok] int32 logical expert ids from the gate.
logcnt (TensorValue) – [num_moe_layers, num_log] int32 replica count per (layer, logical id).
log2phy (TensorValue) – [num_moe_layers, num_log, max_replicas] int32 physical-id table.
layer_idx (TensorValue) – Rank-0 or rank-1 [1] int32 scalar tensor on the same device as router_idx. Rank-0 is reshaped to [1] to match the kernel signature.
num_log (int) – Number of logical experts (comptime).
max_replicas (int) – Maximum replicas per logical expert (comptime).
n_experts_per_tok (int) – Top-K experts per token (comptime). Must be a power of two.
hash_decorrelate (bool) – If True, xor-hash position before the modulo. Defaults to False to preserve legacy routing distribution.

Returns:

[num_tokens, n_experts_per_tok] int32 physical expert ids.

Return type:

Reference: https://github.com/deepseek-ai/DeepSeek-V3/blob/9b4e9788e4a3a731f7567338ed15d3ec549ce03b/inference/model.py#L566

`moe_finalize()`

max.nn.kernels.moe_finalize(down_projs, restore_token_order, router_weight, out_type)

source

Fuses the MoE unpermute gather with the top-k weighted row sum.

Reads each token’s num_experts_per_token expert-permuted rows through restore_token_order, scales each by its router_weight, and sums them, so the [num_tokens, num_experts_per_token, hidden] unpermuted tensor never materializes in HBM.

Parameters:

down_projs (TensorValue) – Expert outputs in expert-permuted (token_expert_order) row order. Shape: [num_tokens * num_experts_per_token, hidden].
restore_token_order (TensorValue) – Maps token-major index i = token * num_experts_per_token + k to its row in down_projs. Shape: [num_tokens * num_experts_per_token], dtype uint32.
router_weight (TensorValue) – Per-(token, expert) routing weight applied before the sum. Shape: [num_tokens, num_experts_per_token].
out_type (DType) – Output dtype.

Returns:

[num_tokens, hidden], dtype out_type.

Return type:

The combined per-token output. Shape

`moe_router_group_limited()`

max.nn.kernels.moe_router_group_limited(expert_scores, expert_bias, n_routed_experts, n_experts_per_tok, n_groups, topk_group, norm_weights, routed_scaling_factor)

source

Routes tokens with the group-limited MoE router.

When n_groups > 1, selects up to topk_group expert groups, then picks n_experts_per_tok experts within those groups (DeepSeek-V3 style). When n_groups == 1, there is only one group, so group selection is skipped and routing uses the dedicated GPU single-group path (mo.moe.single.group.router, implemented as single_group_router in Mojo). In that case topk_group is not used by the kernel.

Parameters:

expert_scores (TensorValue) – The scores for each expert for each token. Shape: [num_tokens, n_routed_experts].
expert_bias (TensorValue) – The bias for each expert. Shape: [n_routed_experts].
n_routed_experts (int) – The total number of experts. Must be divisible by n_groups.
n_experts_per_tok (int) – The number of experts to be selected per token.
n_groups (int) – The total number of expert groups. n_routed_experts must be divisible by this.
topk_group (int) – The maximum number of expert groups that a token will be routed to.
norm_weights (bool) – Whether to normalize the selected expert weights when n_groups > 1. When n_groups == 1, normalization is currently always enabled (norm_weights is treated as True) so behavior matches the graph path that always divided weights by their sum per token.
routed_scaling_factor (float)

Returns:

expert_indices: The indices of the routed experts for each token. Shape: [num_tokens, n_experts_per_tok].
expert_weights: The weights of the routed experts for each token. Shape: [num_tokens, n_experts_per_tok].

Return type:

A tuple of two tensors

The following example routes each of four tokens to two of 16 experts on a GPU. The experts form four groups of four, and each token draws its experts from its two best groups:

import numpy as np
from max.driver import Accelerator, Buffer
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.kernels import moe_router_group_limited

num_tokens, n_routed_experts = 4, 16

with Graph(
    "moe_router_group_limited_example",
    input_types=(
        TensorType(
            DType.float32,
            [num_tokens, n_routed_experts],
            device=DeviceRef.GPU(),
        ),
        TensorType(
            DType.float32, [n_routed_experts], device=DeviceRef.GPU()
        ),
    ),
) as graph:
    expert_scores, expert_bias = (
        value.tensor for value in graph.inputs
    )
    expert_indices, expert_weights = moe_router_group_limited(
        expert_scores,
        expert_bias,
        n_routed_experts=n_routed_experts,
        n_experts_per_tok=2,
        n_groups=4,
        topk_group=2,
        norm_weights=True,
        routed_scaling_factor=2.5,
    )
    graph.output(expert_indices, expert_weights)

rng = np.random.default_rng(0)
# The router takes activated scores, such as sigmoid outputs.
logits = rng.standard_normal((num_tokens, n_routed_experts))
scores = (1 / (1 + np.exp(-logits))).astype(np.float32)
bias = rng.uniform(-0.1, 0.1, n_routed_experts).astype(np.float32)

gpu = Accelerator()
model = InferenceSession(devices=[gpu]).load(graph)
outputs = model.execute(
    Buffer.from_numpy(scores).to(gpu), Buffer.from_numpy(bias).to(gpu)
)
indices, weights = (buffer.to_numpy() for buffer in outputs)
# indices and weights have shape (4, 2). Each token's two weights
# sum to routed_scaling_factor (2.5).

`moe_router_image_bias()`

max.nn.kernels.moe_router_image_bias(expert_scores, expert_bias, image_bias, token_ids, image_token_id, n_routed_experts, n_experts_per_tok, routed_scaling_factor)

source

Routes tokens with a selection bias that depends on the token id.

Selects n_experts_per_tok experts for each token by score plus bias, where the bias is image_bias for tokens whose id is image_token_id and expert_bias for the rest. The weights are the unbiased scores of the selected experts, normalized to sum to 1 (a token whose selected scores are all zero keeps zero weights), then scaled by routed_scaling_factor. Runs one GPU kernel.

Parameters:

expert_scores (TensorValue) – The scores for each expert for each token. Shape: [num_tokens, n_routed_experts].
expert_bias (TensorValue) – The selection bias of the other tokens. Shape: [n_routed_experts].
image_bias (TensorValue) – The selection bias of the image tokens. Shape: [n_routed_experts].
token_ids (TensorValue) – The token ids. Shape: [num_tokens].
image_token_id (int) – The id that marks an image token.
n_routed_experts (int) – The total number of experts.
n_experts_per_tok (int) – The number of experts selected per token.
routed_scaling_factor (float) – The scalar multiplied into every weight.

Returns:

A tuple of the routed experts’ indices and weights, each of shape [num_tokens, n_experts_per_tok].

Raises:

ValueError – If a bias or the token ids do not match the scores’ shape.

Return type:

`moe_router_single_group_eplb()`

max.nn.kernels.moe_router_single_group_eplb(expert_scores, expert_bias, logcnt, log2phy, layer_idx, *, n_routed_experts, n_experts_per_tok, norm_weights, num_log, max_replicas, hash_decorrelate, routed_scaling_factor)

source

Fused single-group MoE router + EPLB log->phy remap.

Replaces the chained moe_router_group_limited (n_groups==1) → moe_eplb_remap for the single-group path. Returns physical expert ids, logical expert ids (kept for the EPLB stats histogram), and routing weights in one launch.

Parameters:

expert_scores (TensorValue)
expert_bias (TensorValue)
logcnt (TensorValue)
log2phy (TensorValue)
layer_idx (TensorValue)
n_routed_experts (int)
n_experts_per_tok (int)
norm_weights (bool)
num_log (int)
max_replicas (int)
hash_decorrelate (bool)
routed_scaling_factor (float)

Return type:

`moe_sink_gate_router()`

max.nn.kernels.moe_sink_gate_router(logits, expert_bias, global_scale, n_routed_experts, n_experts_per_tok, n_shared_experts, route_scale)

source

Fused sigmoid-gate MoE router with always-on sink experts.

Sink experts are gated shared experts, not attention sinks.

Selects the top n_experts_per_tok routed experts by sigmoid(logits) + expert_bias, then softmax-normalizes the log-sigmoid of those experts’ raw logits together with n_shared_experts always-selected sink logits, scaled by route_scale * global_scale. This is the Inkling gate formula, run on the GPU as a single kernel (mo.moe.sink.gate.router, implemented as sink_gate_router in Mojo).

Parameters:

logits (TensorValue) – Raw (pre-sigmoid) gate logits, routed experts followed by sink experts, in any floating dtype. The kernel scores them in float32 and stores the weights in logits.dtype. Shape: [num_tokens, at least n_routed_experts + n_shared_experts]; a wider row’s tail is not read, so a gate weight padded for alignment needs no slice.
expert_bias (TensorValue) – Per-routed-expert selection bias. Shape: [n_routed_experts].
global_scale (TensorValue) – Scalar float32 output-scaling weight. Shape: [1].
n_routed_experts (int) – Total number of routed experts. Must be a positive multiple of the target’s warp width, no greater than 1024: the kernel runs one thread per routed expert.
n_experts_per_tok (int) – Number of routed experts selected per token. Bounded jointly with n_routed_experts by the top-k’s surviving candidates having to fit one warp, so the ceiling falls as the expert count rises: at a warp width of 32, 256 experts admit up to 10 per token and 512 up to 8, while at 64 the same 256 experts admit up to 32.
n_shared_experts (int) – Number of always-selected sink experts. n_experts_per_tok + n_shared_experts must be a power of two no greater than the target’s warp width, since the two are jointly normalized by a single warp-level reduction.
route_scale (float) – Scalar output-scaling factor.

Returns:

expert_ids: Selected routed-expert indices. Shape: [num_tokens, n_experts_per_tok].
expert_weights: Routing weight per selected routed expert. Shape: [num_tokens, n_experts_per_tok].
sink_weights: Routing weight per sink expert. Shape: [num_tokens, n_shared_experts].

Return type:

A tuple of three tensors

Raises:

ValueError – If the routing shape is one the kernel cannot compile.

`msa_sparse_attention_ragged()`

max.nn.kernels.msa_sparse_attention_ragged(kv_params, input, input_row_offsets, cache_row_offsets, total_context_length, kv_collection, layer_idx, block_indices, *, group, topk, sparse_block_size, scale)

source

Computes MiniMax-M3 block-sparse attention over the main paged KV cache.

Gathers topk 128-token KV blocks per (kv head, query token) using the block ids produced by msa_sparse_indexer(), then runs SM100 block-sparse MHA. The main KV cache is BF16, native FP8 e4m3, or NVFP4; input (the query) must match its dtype, or be FP8 e4m3 for NVFP4, and the output is always BF16. head_dim is 128. The op selects the prefill or decode kernel at runtime from the main KV cache’s max_seq_length, so the same call serves both paths.

Parameters:

kv_params (KVCacheParams) – Key-value cache parameters for the main KV cache.
input (TensorValue) – Query tensor [total_q, n_heads, head_dim] (prefill) or [batch, n_heads, head_dim] (decode); dtype matches the KV cache (BF16 or FP8 e4m3).
input_row_offsets (TensorValue) – Ragged query offsets [batch + 1] uint32.
cache_row_offsets (TensorValue) – Ragged valid-cache offsets [batch + 1] uint32.
total_context_length (TensorValue) – Total padded cache length for the batch, CPU scalar [1] uint32.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Main paged KV cache (BF16 or FP8 e4m3, no scales).
layer_idx (TensorValue) – Layer index, uint32, on CPU.
block_indices (TensorValue) – Selected block ids. Prefill: [n_kv_heads, total_q, topk]; decode: [n_kv_heads, batch, topk]. int32.
group (int) – Query heads per kv-head (n_heads // n_kv_heads).
topk (int) – Number of gathered KV blocks per token.
sparse_block_size (int) – KV block size in tokens; the model’s sparse_attention_config.sparse_block_size. Must equal the KV cache page size and the kernel’s BN.
scale (float) – QK scale.

Returns:

Output tensor [total_q, n_heads, head_dim] (prefill) or [batch, n_heads, head_dim] (decode), BF16.

Return type:

`msa_sparse_attention_ragged_mxfp6()`

max.nn.kernels.msa_sparse_attention_ragged_mxfp6(kv_params, input, input_row_offsets, cache_row_offsets, total_context_length, kv_collection, layer_idx, block_indices, *, group, topk, scale, fp6_format='e2m3')

source

Computes MiniMax-M3 block-sparse attention, emitting packed MXFP6 + scales.

AMD (gfx950) variant of msa_sparse_attention_ragged() whose output is the o_proj-ready MXFP8 activation instead of BF16: quantized data [num_rows, n_heads, head_dim] in float8_e4m3fn plus E8M0 block scales [num_rows, n_heads * head_dim // 32] – the same pair quantize_dynamic_block_scaled() produces from the BF16 output, so it feeds dynamic_block_scaled_matmul_amd() directly and the separate quantize dispatch is skipped. Bit-identical to that unfused pair; on split-K decode shapes the quantize fuses into the reduce and saves a dispatch.

Parameters:

kv_params (KVCacheParams) – Key-value cache parameters for the main KV cache.
input (TensorValue) – Query tensor [total_q, n_heads, head_dim] (prefill) or [batch, n_heads, head_dim] (decode); dtype matches the KV cache (BF16 or FP8 e4m3).
input_row_offsets (TensorValue) – Ragged query offsets [batch + 1] uint32.
cache_row_offsets (TensorValue) – Ragged valid-cache offsets [batch + 1] uint32.
total_context_length (TensorValue) – Total padded cache length for the batch, CPU scalar [1] uint32.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Main paged KV cache (BF16 or FP8 e4m3, no scales).
layer_idx (TensorValue) – Layer index, uint32, on CPU.
block_indices (TensorValue) – Selected block ids. Prefill: [n_kv_heads, total_q, topk]; decode: [n_kv_heads, batch, topk]. int32.
group (int) – Query heads per kv-head (n_heads // n_kv_heads).
topk (int) – Number of gathered KV blocks per token.
scale (float) – QK scale.
fp6_format (str)

Returns:

The quantized attention output [total_q, n_heads, head_dim] float8_e4m3fn and its E8M0 block scales [total_q, n_heads * head_dim // 32].

Return type:

`msa_sparse_attention_ragged_mxfp8()`

max.nn.kernels.msa_sparse_attention_ragged_mxfp8(kv_params, input, input_row_offsets, cache_row_offsets, total_context_length, kv_collection, layer_idx, block_indices, *, group, topk, sparse_block_size, scale)

source

Computes MiniMax-M3 block-sparse attention, emitting MXFP8 + scales.

Parameters:

kv_params (KVCacheParams) – Key-value cache parameters for the main KV cache.
input (TensorValue) – Query tensor [total_q, n_heads, head_dim] (prefill) or [batch, n_heads, head_dim] (decode); dtype matches the KV cache (BF16 or FP8 e4m3), or is FP8 e4m3 for NVFP4.
input_row_offsets (TensorValue) – Ragged query offsets [batch + 1] uint32.
cache_row_offsets (TensorValue) – Ragged valid-cache offsets [batch + 1] uint32.
total_context_length (TensorValue) – Total padded cache length for the batch, CPU scalar [1] uint32.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Main paged KV cache: BF16 or FP8 e4m3 without scales, or NVFP4 with its E4M3 scales.
layer_idx (TensorValue) – Layer index, uint32, on CPU.
block_indices (TensorValue) – Selected block ids. Prefill: [n_kv_heads, total_q, topk]; decode: [n_kv_heads, batch, topk]. int32.
group (int) – Query heads per kv-head (n_heads // n_kv_heads).
topk (int) – Number of gathered KV blocks per token.
sparse_block_size (int) – KV block size in tokens; the model’s sparse_attention_config.sparse_block_size. Must equal the KV cache page size and the kernel’s BN.
scale (float) – QK scale.

Returns:

The quantized attention output [total_q, n_heads, head_dim] float8_e4m3fn and its E8M0 block scales [total_q, n_heads * head_dim // 32].

Return type:

`msa_sparse_indexer()`

max.nn.kernels.msa_sparse_indexer(kv_params, index_q, input_row_offsets, prefix_lens, index_kv_collection, layer_idx, score_scratch, *, num_index_heads, idx_head_dim, block_size, topk, init_blocks, local_blocks, scale)

source

Selects the top-k key blocks per token for MiniMax-M3 sparse attention.

Computes per-block QK scores against the index-K paged cache (BF16 or scale-free FP8 e4m3) and returns, for each (index head, query token), the ids of the topk highest-scoring 128-token blocks. The forward sparse-attention op consumes these block ids to gather a sparse KV band. The op selects the prefill or decode kernel at runtime from the index-K cache’s max_seq_length, so the same call serves both paths.

Parameters:

kv_params (KVCacheParams) – Key-value cache parameters for the index-K cache.
index_q (TensorValue) – Index queries. Prefill: [total_q, num_index_heads, idx_head_dim]; decode: [batch, num_index_heads, idx_head_dim]. BF16.
input_row_offsets (TensorValue) – Ragged query offsets [batch + 1] uint32 (used on the prefill path; on decode it is [0, 1, ..., batch]).
prefix_lens (TensorValue) – Per-row cached-key count [batch] (the index-K cache_lengths). uint32.
index_kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – Paged index-K cache (BF16 or scale-free e4m3, no scales).
layer_idx (TensorValue) – Layer index, uint32, on CPU.
score_scratch (BufferValue) – Persistent FP32 scratch buffer for decode scoring, shape [num_index_heads, max_rows, MAX_NUM_BLOCKS]. max_rows must cover total_q for every decode step served: max_batch for single-token decode, and max_batch * (num_draft_tokens + 1) for the multi-token (MTP / speculative) routes, which write at the ragged row input_row_offsets[b] + t. It is persistent because decode runs inside the graph-capture region, so the kernel cannot allocate a correctly sized buffer per step; a scratch too narrow for a multi-token route makes that step fall back to the prefill indexer, while single-token decode raises.
num_index_heads (int) – Number of index (query) heads.
idx_head_dim (int) – Index head dimension.
block_size (int) – KV block size in tokens (== page size).
topk (int) – Number of blocks to select per token.
init_blocks (int) – Always-keep leading blocks.
local_blocks (int) – Always-keep trailing/local blocks.
scale (float) – QK scale.

Returns:

[num_index_heads, total_q, topk]; decode: [num_index_heads, batch, topk].

Return type:

Block indices, int32, -1-padded. Prefill

`mtp_eh_norm()`

max.nn.kernels.mtp_eh_norm(embed, prev_hidden, enorm_weight, hnorm_weight, epsilon, block_threads=256)

source

Normalizes both inputs of an MTP draft layer into one projection input.

A multi-token-prediction draft predicts a token two ahead from the embedding of the token the target just produced and the target’s hidden state. Each is RMS-normalized with its own weight, and the pair is projected back to one hidden width. This returns that projection’s input in a single pass.

Normalization matches rms_norm() in its Llama-style configuration – weight_offset=0 and multiply_before_cast=False. Do not use this with the Gemma-style configuration.

Parameters:

embed (TensorValue) – Token embeddings [num_tokens, hidden_size].
prev_hidden (TensorValue) – Target hidden states, same shape and dtype as embed.
enorm_weight (TensorValue) – Embedding norm weight [hidden_size].
hnorm_weight (TensorValue) – Hidden-state norm weight [hidden_size].
epsilon (float) – Added inside the square root.
block_threads (int) – Threads per block; must be a whole number of warps.

Returns:

[num_tokens, 2 * hidden_size] with the normalized embedding in the leading columns and the normalized hidden state in the trailing ones.

Raises:

ValueError – If the inputs disagree in shape, dtype or device, if the hidden size is not statically known, if the dtype is float8, or if block_threads is not a positive multiple of 32 no greater than 1024.

Return type:

`mxfp4_dequant()`

max.nn.kernels.mxfp4_dequant(packed_weights, scales, out_type=bfloat16)

source

Dequantizes MXFP4 packed weights to BF16 or FP8 on GPU.

Supports rank 2 [N, K//2] and rank 3 [E, N, K//2] inputs. For rank 3, leading dims are flattened to 2D, dequantized, and reshaped back.

Parameters:

packed_weights (TensorValue) – Packed weights in uint8 (2 FP4 values per byte). Shape [N, K//2] or [E, N, K//2].
scales (TensorValue) – Block scales in float8_e8m0fnu. Shape [N, K//32] or [E, N, K//32].
out_type (DType) – Output dtype (bfloat16 or float8_e4m3fn).

Returns:

Dequantized tensor [N, K] or [E, N, K] in out_type.

Return type:

`mxfp6_dequant()`

max.nn.kernels.mxfp6_dequant(packed_weights, scales, fp6_format='e2m3', out_type=bfloat16)

source

Dequantizes MXFP6 packed weights to BF16 or FP8 on GPU.

The FP6 sibling of mxfp4_dequant(). Supports rank 2 [N, K * 3 // 4] and rank 3 [E, N, K * 3 // 4] inputs; for rank 3 the leading dims are flattened, dequantized, and reshaped back.

Parameters:

packed_weights (TensorValue) – Packed FP6 bytes, four codes per three bytes.
scales (TensorValue) – Block scales in float8_e8m0fnu, [..., N, K // 32].
fp6_format (str) – The FP6 element encoding, "e2m3" or "e3m2".
out_type (DType) – Output dtype (bfloat16 or float8_e4m3fn).

Returns:

The dequantized tensor, [N, K] or [E, N, K].

Return type:

`needs_fp8_fnuz_conversion()`

max.nn.kernels.needs_fp8_fnuz_conversion()

source

Checks if FP8 E4M3FN to FNUZ conversion is needed for AMD GPUs.

Returns:: True if running on AMD GPU with CDNA3 architecture, False otherwise.
Return type:: bool

`normalize_e4m3fn_to_e4m3fnuz()`

max.nn.kernels.normalize_e4m3fn_to_e4m3fnuz(weight, weight_scale)

source

Converts E4M3FN weights to E4M3FNUZ format for AMD GPUs.

This conversion is necessary because AMD GPUs use the E4M3FNUZ format while NVIDIA GPUs use E4M3FN. The key differences are:

The bit pattern 10000000 (-128) represents zero in E4M3FN but NaN in E4M3FNUZ
For the same bit representation, E4M3FNUZ values are half of E4M3FN values

Parameters:

weight (TensorValue) – The weight tensor in E4M3FN format.
weight_scale (TensorValue) – The weight scale factor.

Returns:

Tuple of (converted_weight, adjusted_weight_scale, adjusted_input_scale).

Return type:

`nvfp4_mla_prefill_convert_bytes()`

max.nn.kernels.nvfp4_mla_prefill_convert_bytes(max_query_tokens, top_k, row_dim, page_size)

source

Returns the most one sparse MLA prefill call over an NVFP4 cache allocates, on one device, to convert its context to FP8 once.

The call allocates the FP8 copy of at most max_rows context rows, and of at most ratio rows per query token; each query token’s top_k indices rewritten into that copy as int32; and a uint32 per cache block mapping it to its page of the copy. A step’s query tokens are rows of its context, so they too are at most max_rows. Cache rows are addressed as int32, which bounds the blocks at 2**31 / page_size.

Parameters:

max_query_tokens (int) – Most query tokens in one prefill step.
top_k (int) – Context rows each query token attends to.
row_dim (int) – Elements of one latent cache row (latent plus rope).
page_size (int) – Tokens per KV cache page.

Returns:

The bytes, or 0 when the conversion is off.

Return type:

int

`nvfp4_mla_prefill_convert_max_rows()`

max.nn.kernels.nvfp4_mla_prefill_convert_max_rows()

source

Returns the most rows the NVFP4 prefill conversion converts.

Return type:: int

`nvfp4_mla_prefill_convert_ratio()`

max.nn.kernels.nvfp4_mla_prefill_convert_ratio()

source

Returns the rows-per-query-token bound of the NVFP4 prefill conversion.

Return type:: int

`preshuffle_block_scaled_b()`

max.nn.kernels.preshuffle_block_scaled_b(weight, lane_bytes=16)

source

Permutes MX expert weights into the AMD preb kernel’s B layout.

A pure byte permutation of constant weights, so the graph compiler hoists it into model init. With 16-byte lanes (MXFP4, MXFP8) each [N, K_BYTES] expert becomes (N0, K0, KLane=4, NLane=16, KPack=16), the b_5d_grouped_layout the kernel reads with coalesced loads. With MXFP6_LANE_BYTES, no power-of-two load covers a lane, so each MFMA tile stores the first 16 bytes of every lane fragment as one plane and the last 8 as a second, each ordered (KLane, NLane, byte). Lanes are counted in bytes, so a wider element type permutes as its bytes would: a bf16 weight lands in the B layout of smallm_streaming_matmul().

Parameters:

weight (TensorValue) – Row-major expert weights [E, N, K] as the checkpoint stores them, with any graph-side reordering of N already applied.
lane_bytes (int) – Bytes one lane feeds the MFMA: 16, or MXFP6_LANE_BYTES for MXFP6.

Returns:

The permuted weights, same shape and dtype.

Raises:

ValueError – If the weight is not rank 3, N is not a multiple of 16, or a row is not a whole number of MFMA K tiles.

Return type:

`preshuffle_block_scaled_b_scales()`

max.nn.kernels.preshuffle_block_scaled_b_scales(scales)

source

Permutes E8M0 expert weight scales into the AMD preb kernel’s layout.

Each [MN, K_SCALES] expert becomes (MN_block, K_block, K_lane=4, MN_lane=16, K_pack=2, MN_pack=2): every i32 cell holds the 2x2 (mn_pack, k_pack) bytes one lane loads at once, in the order the kernel’s OPSEL byte selector reads them (Shuffler.scale_4d_byte_off). Hoisted into model init like preshuffle_block_scaled_b().

Parameters:: scales (TensorValue) – Row-major E8M0 scales [E, MN, K_SCALES].
Returns:: The permuted scales, same shape and dtype.
Raises:: ValueError – If the scales are not rank 3, MN is not a multiple of 32, or K_SCALES is not a multiple of 8.
Return type:: TensorValue

`quantize_dynamic_block_scaled()`

max.nn.kernels.quantize_dynamic_block_scaled(input, tensor_sf=1.0, sf_vector_size=16, scales_type=float8_e4m3fn, out_type=uint8)

source

Dynamically quantize a bf16 tensor to NVFP4 or MXFP8 with block scales.

Both formats go through the mo.quantize.dynamic.block.scaled op; the format is selected from out_type:

NVFP4 (out_type=uint8): two e2m1 values packed per output byte, with float8_e4m3fn (NVFP4) or float8_e8m0fnu (MXFP4) scales and sf_vector_size 16 or 32.
MXFP8 (out_type=float8_e4m3fn): unpacked float8_e4m3fn data with float8_e8m0fnu scales and sf_vector_size=32. MXFP8 uses pure block scaling, so tensor_sf defaults to 1.0 (identity).

Parameters:

input (TensorValue) – The input tensor to quantize. Shape: [seq_len, hidden_size].
tensor_sf (TensorValue | float) – The tensor-wise scale factor (inverted as per the quantization kernel requirement for NVFP4; identity for MXFP8).
sf_vector_size (int) – The block size for the scaling factors. 16 for NVFP4, 32 for MXFP4/MXFP8.
scales_type (DType) – The type of the scales tensor. float8_e4m3fn for NVFP4, float8_e8m0fnu for MXFP4/MXFP8.
out_type (DType) – The type of the output tensor. uint8 for packed FP4 or float8_e4m3fn for MXFP8.

Returns:

rank-5 interleaved on NVIDIA SM100, rank-2 [M, K // sf_vector_size] otherwise.

Return type:

The quantized tensor and scales. Scales layout depends on hardware

`quantize_dynamic_block_scaled_mxfp4()`

max.nn.kernels.quantize_dynamic_block_scaled_mxfp4(input, scales_type=float8_e8m0fnu, out_type=uint8)

source

Dynamically quantize the input tensor to fp4-e2m1fn.

Parameters:

input (TensorValue) – The input tensor to quantize. Shape: [seq_len, hidden_size]
out_type (DType) – The type of the output tensor.
scales_type (DType) – The type of the scales tensor.

Returns:

The quantized tensor in [seq_len, hidden_size // 2] layout and the scales in [seq_len, hidden_size // 32] layout.

Return type:

`quantize_dynamic_block_scaled_mxfp6()`

max.nn.kernels.quantize_dynamic_block_scaled_mxfp6(input, fp6_format='e2m3', scales_type=float8_e8m0fnu, out_type=uint8)

source

Dynamically quantizes the input tensor to MXFP6.

Parameters:

input (TensorValue) – The tensor to quantize, [seq_len, hidden_size] bf16.
fp6_format (str) – The FP6 element encoding, "e2m3" or "e3m2".
scales_type (DType) – The dtype of the scales tensor.
out_type (DType) – The dtype of the packed output.

Returns:

The packed tensor in [seq_len, hidden_size * 3 // 4] and the scales in [seq_len, hidden_size // 32].

Return type:

`quantize_dynamic_scaled_float8()`

max.nn.kernels.quantize_dynamic_scaled_float8(input, input_scale_spec, weight_scale_spec, scale_ub=1200.0, group_size_or_per_token=-1, out_type=float8_e4m3fn, scales_type=bfloat16, amax_floor=0.0, row_offsets=None)

source

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.
row_offsets (TensorValue | None) – Optional grouped-matmul row prefix sum, whose last entry is the number of rows that carry data. Supply it when input is a worst-case-sized buffer that the producer only partly filled, such as the EP dispatch receive buffer; rows past that count are left untouched rather than quantized. The quantized values and scales of the rows below it are identical either way.
out_type (DType) – The type of the output tensor.
scales_type (DType) – The type of the scales tensor.
amax_floor (float) – Lower bound applied to a group’s max-abs before the scale is derived. QAT recipes that floor the activation amax (DeepSeek-V4’s act_quant uses 1e-4) need it to stay on the scale grid the checkpoint was trained against. Supported only when scales_type is float8_e8m0fnu; 0.0 (the default) leaves the scale unchanged.
input_scale_spec (InputScaleSpec)
weight_scale_spec (WeightScaleSpec)

Returns:

The quantized tensor and the scales.

Return type:

`quantize_static_scaled_float8()`

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

source

Quantizes a rank-2 tensor to float8 using a static per-tensor scale.

Parameters:

x (TensorValue) – Input tensor to quantize. Must be rank 2 with dtype float16, bfloat16, or float32.
scale (TensorValue) – Scalar scale factor (shape [] or [1]) residing on CPU.
scale_is_inverted (bool) – When True (default), scale is interpreted as 1 / max_val (inverted). When False, it is the raw absolute-max scale.
out_type (DType) – Output dtype. Defaults to DType.float8_e4m3fn.

Returns:

A quantized TensorValue with shape equal to x and dtype out_type.

Raises:

ValueError – If scale is not a scalar, x is not rank 2, x dtype is unsupported, or scale is not on CPU.

Return type:

`quantize_tensor_dynamic_scaled_float8()`

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

source

Quantizes a rank-2 tensor to float8 using a dynamic per-tensor scale.

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.
input_scale_spec (InputScaleSpec)
weight_scale_spec (WeightScaleSpec)

Returns:

The quantized tensor and the scales.

Return type:

`repack_gguf_quantized_weights()`

max.nn.kernels.repack_gguf_quantized_weights(weight, quantization_encoding)

source

Repacks GGUF quantized weights for the given encoding.

Parameters:

weight (TensorValue) – The quantized weight tensor to repack.
quantization_encoding (QuantizationEncoding) – The quantization encoding to repack for.

Returns:

A TensorValue representing the repacked weights.

Return type:

`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, multiply_before_cast=True, per_head_norm=True)

source

Applies RMSNorm to the new entries in the KV cache.

When per_head_norm is True (the default), RMSNorm is applied separately to each head. In this mode, gamma should have size [head_dim] and normalization occurs across the head_dim dimensions within each head.

When per_head_norm is False, RMSNorm is applied per token across all heads. In this mode, gamma should have size [n_kv_heads * head_dim] and normalization occurs across all dimensions for each token.

The size of the gamma tensor determines how many dimensions will be normalized. If gamma’s size doesn’t match the expected size based on the per_head_norm setting, rms_norm_cols must be explicitly specified to confirm the intention to normalize only a subset of dimensions.

The KV cache collection itself isn’t aware of the new cache entries until the cache length increment, which happens after the model forward, so input_row_offsets does this bookkeeping.

Parameters:

kv_params (KVCacheParams) – The KV cache parameters.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache holding the entries to normalize.
gamma (TensorValue) – The RMSNorm weight.
epsilon (float | floating[Any]) – The epsilon added inside the normalization.
layer_idx (TensorValue) – The index of the layer being normalized.
total_seq_len (Dim) – The total sequence length of the ragged batch.
input_row_offsets (TensorValue) – The ragged offsets delimiting the new entries.
weight_offset (float | floating[Any]) – The offset added to gamma before the multiply.
rms_norm_cols (int | None) – The number of columns to normalize. Required when gamma’s size doesn’t match the expected size, to confirm the intention to normalize only a subset of dimensions.
multiply_before_cast (bool) – Whether to multiply by gamma before casting back to the cache dtype.
per_head_norm (bool) – Whether to normalize each head separately. Defaults to True.

Return type:

None

`rms_norm_quantize_dynamic_scaled_float8()`

max.nn.kernels.rms_norm_quantize_dynamic_scaled_float8(input, gamma, epsilon, input_scale_spec, weight_scale_spec, *, weight_offset=0.0, multiply_before_cast=False, scale_ub=1200.0, group_size_or_per_token=-1, out_type=float8_e4m3fn, scales_type=float32)

source

RMS-normalizes input and quantizes the result to FP8 with dynamic group (blockwise) scales in a single kernel.

Equivalent to rms_norm() followed by quantize_dynamic_scaled_float8() with the same arguments, including the [K // group_size, M_padded] scale layout, without materializing the normalized activation. Group scales only: a per-token group (group_size == K) is rejected, that case fuses through the graph compiler’s rms_norm + quantize composite instead.

Parameters:

input (TensorValue) – The [M, K] activation to normalize and quantize.
gamma (TensorValue) – The [K] RMSNorm weight, already in input’s dtype and on its device.
epsilon (float) – The RMSNorm epsilon.
input_scale_spec (InputScaleSpec) – Activation scale spec; supplies the group size when group_size_or_per_token is -1.
weight_scale_spec (WeightScaleSpec) – Weight scale spec of the consuming matmul.
weight_offset (float) – Constant added to gamma (Gemma-style norms).
multiply_before_cast (bool) – RMSNorm gamma-multiply order; see rms_norm().
scale_ub (float) – Upper bound of the dynamic scale.
group_size_or_per_token (int) – Columns per scale, or -1 to take the activation block size from input_scale_spec.
out_type (DType) – FP8 output dtype.
scales_type (DType) – Scale dtype. Defaults to float32, which the blockwise-FP8 matmuls require.

Returns:

The quantized [M, K] tensor and its [K // group_size, M_padded] scales.

Return type:

`rms_norm_quantize_mxfp8_q_kv()`

max.nn.kernels.rms_norm_quantize_mxfp8_q_kv(x, q_weight, kv_weight, epsilon, q_cols)

source

RMS-normalizes the q and kv parts of each row and quantizes q to MXFP8.

Each row of x is q_cols columns of q then the rest as kv. One kernel computes what rms_norm() on each part, then quantize_dynamic_block_scaled() on the normalized q, would, without writing the normalized q out. The normalized rows are rounded to bfloat16 where rms_norm rounds them. The two paths differ only in the order of each row’s sum of squares. Runs on NVIDIA SM100 and SM120.

Parameters:

x (TensorValue) – The bfloat16 input. Shape: [M, q_cols + kv_cols].
q_weight (TensorValue) – The bfloat16 q weight. Shape: [q_cols].
kv_weight (TensorValue) – The bfloat16 kv weight. Shape: [kv_cols].
epsilon (float) – Added to each row’s mean square.
q_cols (int) – The width of the q part. q_cols and kv_cols must each be a multiple of 256.

Returns:

The float8_e4m3fn q [M, q_cols], its float8_e8m0fnu scales in the rank-5 interleaved layout, and the bfloat16 normalized kv [M, kv_cols].

Raises:

ValueError – If the GPU is not SM100 or SM120, or the shapes and dtypes are not as above.

Return type:

`rms_norm_value_cache()`

max.nn.kernels.rms_norm_value_cache(kv_params, kv_collection, gamma, epsilon, layer_idx, total_seq_len, input_row_offsets, weight_offset, rms_norm_cols=None, multiply_before_cast=True, per_head_norm=True)

source

Applies RMSNorm in place to the _new_ entries in the value cache. Semantics match rms_norm_key_cache(), but updates the value tensor for the layer instead of the key tensor.

Parameters:

kv_params (KVCacheParams)
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])
gamma (TensorValue)
epsilon (float | floating[Any])
layer_idx (TensorValue)
total_seq_len (Dim)
input_row_offsets (TensorValue)
weight_offset (float | floating[Any])
rms_norm_cols (int | None)
multiply_before_cast (bool)
per_head_norm (bool)

Return type:

None

`rope_quantize_dynamic_block_scaled()`

max.nn.kernels.rope_quantize_dynamic_block_scaled(x, freqs, head_dim, inverse=False, tensor_sf=1.0, sf_vector_size=16, scales_type=float8_e4m3fn, out_type=uint8)

source

Rotates each head’s trailing columns by RoPE and quantizes with block scales.

Each row of x holds whole heads of head_dim columns. One kernel rotates the last rope_dim columns of every head as interleaved (re, im) pairs by the row’s freqs, rounds them to bfloat16 and quantizes the row as quantize_dynamic_block_scaled() would, without writing the rotated input out. The rotation is rope_ragged()’s interleaved arithmetic. The formats and the remaining arguments match quantize_dynamic_block_scaled(). Runs on NVIDIA SM100 and SM120.

Parameters:

x (TensorValue) – The bfloat16 input. Shape: [M, N] with N a multiple of head_dim.
freqs (TensorValue) – The float32 (cos, sin) pairs of each row. Shape: [M, rope_dim // 2, 2] or [M, rope_dim].
head_dim (int) – Columns per head.
inverse (bool) – Rotate by the conjugate of each frequency, which undoes a forward rotation.
tensor_sf (TensorValue | float) – The tensor-wise scale factor (inverted for NVFP4; identity for MXFP8).
sf_vector_size (int) – The block size for the scaling factors. 16 for NVFP4, 32 for MXFP4/MXFP8.
scales_type (DType) – The type of the scales tensor. float8_e4m3fn for NVFP4, float8_e8m0fnu for MXFP4/MXFP8.
out_type (DType) – The type of the output tensor. uint8 for packed FP4 or float8_e4m3fn for MXFP8.

Returns:

The quantized [M, N] input ([M, N // 2] packed FP4) and its rank-5 interleaved scales.

Raises:

ValueError – If the GPU is not SM100 or SM120, x is not a rank-2 bfloat16 tensor of whole heads, freqs does not hold one float32 row per row of x with a static width of whole 8-column chunks, or the format arguments do not name a supported format.

Return type:

`rope_ragged()`

max.nn.kernels.rope_ragged(input, input_row_offsets, start_pos, freqs_cis, *, interleaved=True, rope_first=False, output_dtype=None)

source

Applies RoPE to ragged input using the standard rope kernel.

When freqs_cis is narrower than input’s head dimension, only that many columns of each head are rotated and the rest pass through. Those rotated columns are the trailing ones by default; set rope_first to rotate the leading ones instead.

Parameters:

input (TensorValue)
input_row_offsets (TensorValue)
start_pos (TensorValue)
freqs_cis (TensorValue)
interleaved (bool)
rope_first (bool)
output_dtype (DType | None)

Return type:

`rope_ragged_with_position_ids()`

max.nn.kernels.rope_ragged_with_position_ids(input, freqs_cis, position_ids, *, mrope_section=None, interleaved=True)

source

Applies RoPE using explicit position_ids (no KV cache coupling).

Parameters:

input (TensorValue)
freqs_cis (TensorValue)
position_ids (TensorValue)
mrope_section (list[int] | None)
interleaved (bool)

Return type:

`rope_split_store_ragged()`

max.nn.kernels.rope_split_store_ragged(kv_params, qkv, input_row_offsets, freqs_cis, kv_collection, layer_idx, n_heads, interleaved=True, position_ids=None, mrope_section=None, fuse=True, q_out_dtype=None, q_norm_weight=None, k_norm_weight=None, rms_norm_eps=None, k_eq_v=False)

source

Applies RoPE to Q and K from a flat QKV buffer and stores K/V to the cache.

Reads from a flat QKV matmul output, applies RoPE to the Q and K regions, stores K/V to the paged KV cache, and writes the roped Q to the output.

Parameters:

kv_params (KVCacheParams) – The KV cache parameters.
qkv (TensorValue) – The flat QKV matmul output, [total_seq_len, q_dim + k_dim + v_dim].
input_row_offsets (TensorValue) – The ragged offsets, [batch_size + 1].
freqs_cis (TensorValue) – The RoPE frequencies, [max_seq_len, head_dim].
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache.
layer_idx (TensorValue) – The layer index.
n_heads (int) – The number of query attention heads.
interleaved (bool) – Whether freqs_cis uses interleaved (re, im) format.
position_ids (TensorValue | None) – The optional ragged 2D array of position IDs. If None, defaults to cache_length + token_idx for each token. When num_sections > 1, mrope_section must be provided. Shape: [num_sections, total_seq_len].
mrope_section (list[int] | None) – The optional list of ints indicating the section of the head_dim to apply RoPE to. Must be used with position_ids.
fuse (bool) – If True (the default), emits a single fused custom op. If False, emits separate split, rope, and store ops for testing graph compiler fusion.
q_out_dtype (DType | None) – The dtype for the roped Q output. Defaults to qkv.dtype.
q_norm_weight (TensorValue | None) – The optional per-head RMSNorm gamma [head_dim] for Q. When given (with k_norm_weight and rms_norm_eps), the per-head Q/K/V RMS-norm is fused into the op (Q/K use their gammas, V is a bare norm), removing the separate norm ops. Mutually exclusive with position_ids.
k_norm_weight (TensorValue | None) – The per-head RMSNorm gamma [head_dim] for K (see q_norm_weight).
rms_norm_eps (float | None) – The epsilon for the fused qk-norm; required when q_norm_weight is set.
k_eq_v (bool) – When True (only valid with q_norm_weight), V has no own projection and reuses K’s: qkv is [q|k] (no V region) and the kernel reads the K head for both the K and V stores, sharing the norm reduction. When False (the default), qkv is [q|k|v].

Returns:

The roped Q output, [total_seq_len, n_heads * head_dim].

Return type:

`row_mean_of_squares()`

max.nn.kernels.row_mean_of_squares(x)

source

Computes the per-row mean of squares over the last axis.

For an input x flattened to [M, N] over its last axis, computes out[m] = sum_n(float32(x[m, n]) ** 2) / N. The square and accumulation always run in float32 regardless of the input dtype, and the result is always float32. The output preserves the leading axes with a trailing size-1 reduction axis, matching ops.mean(x * x, axis=-1).

This is a fused, single-pass replacement for ops.mean(x * x, axis=-1) used in QK-RMSNorm-style variance computations. The generic reduce path over-provisions the grid for small M (decode); this op launches exactly one block per row.

Parameters:: x (TensorValue) – The input tensor. Reduction runs over the last axis. Accepts bfloat16 or float32 (any rank >= 1).
Returns:: A float32 TensorValue whose shape matches x with the last axis replaced by 1.
Raises:: ValueError – If x dtype is not bfloat16 or float32.
Return type:: TensorValue

`row_mean_of_squares_qk()`

max.nn.kernels.row_mean_of_squares_qk(q, k)

source

Fused per-row mean of squares for two operands q and k.

For q of shape [M, Nq] and k of shape [M, Nk] (sharing the leading rows dimension but with possibly different column counts), computes a [M, 2] float32 result where column 0 is mean_n(q[m, n] ** 2) and column 1 is mean_n(k[m, n] ** 2). The square and accumulation always run in float32 regardless of input dtype.

This is a single-kernel fusion of two row_mean_of_squares() calls plus a concat, equivalent to ops.concat([row_mean_of_squares(q), row_mean_of_squares(k)], axis=-1) but with one launch instead of two plus the concat. It is used for the cross-head QK-RMSNorm variance statistics in tensor-parallel attention.

Parameters:

q (TensorValue) – The Q projection, [M, Nq]. Reduction runs over the last axis.
k (TensorValue) – The K projection, [M, Nk]. Must share q’s rows dim and dtype.

Returns:

A float32 TensorValue of shape [M, 2].

Raises:

ValueError – If q/k are not rank 2, have mismatched dtypes or rows, or use a dtype other than bfloat16 or float32.

Return type:

`scatter_nd_skip_oob_indices()`

max.nn.kernels.scatter_nd_skip_oob_indices(input, updates, indices)

source

Creates a new symbolic tensor where the updates are scattered into input at specified indices.

This differs from scatter_nd in that it handles oob indices by skipping the update for that index. Oob indices are those which fall outside of the range [-dim, dim).

Parameters:

input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to write elements to.
updates (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – A symbolic tensor of elements to write to input.
indices (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – A tensor of indices specifying where to write updates. Shape should be [num_updates, rank] for full indexing or [num_updates, k] for partial indexing where k < rank.

Returns:

A new symbolic tensor representing the result of the scatter_nd operation.

Return type:

`scatter_set_constant()`

max.nn.kernels.scatter_set_constant(data, indices, fill_val)

source

Scatters values into a tensor at specified indices.

Parameters:

data (BufferValue | HasBufferValue)
indices (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray)
fill_val (float)

Return type:

None

`sgmv_kernel()`

max.nn.kernels.sgmv_kernel(input, lora, lora_ids, lora_ranks, input_row_offsets, max_lora_seq_len, lora_end_idx=None, bias=None)

source

Performs the SGMV kernel for LoRA. This is LoRA agnostic, meaning that we can perform LoRA A or B from this kernel call.

Parameters:

input (TensorValue) – The input tensor.
lora (TensorValue) – The LoRA tensor.
lora_ids (TensorValue) – Ids of the LoRAs used for each sequence
lora_ranks (TensorValue) – The ranks of the LoRAs in the batch.
input_row_offsets (TensorValue) – The sequence offsets that use LoRA
max_lora_seq_len (int) – The maximum sequence length of any given LoRA in the batch
bias (TensorValue | None) – The LoRA bias
lora_end_idx (TensorValue | None)

Raises:

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

`sgmv_lora_kernel()`

max.nn.kernels.sgmv_lora_kernel(input, lora_a, lora_b, lora_ids, lora_ranks, grouped_row_offsets, lora_end_idx, max_lora_seq_len, bias=None)

source

Computes the SGMV LoRA kernel for some number of LoRAs A and B given the input.

out = Wx + xAB

SGMV can be explained by two independent kernels:: shrink -> shrinks high-dimensional tensor to low-rank tensor

expand -> expands low-rank tensor to high-dimensional tensor

where v = [0, …] and y = (some output tensor)

SGMV-shrink:: v += xA
SGMV-expand:: y += vB

Parameters:

input (TensorValue) – The input tensor
lora_a (TensorValue) – The LoRA tensor for A
lora_b (TensorValue) – The LoRA tensor for B
lora_ids (TensorValue) – Ids of the LoRAs used for each sequence
lora_ranks (TensorValue) – The ranks of the LoRAs in the batch.
grouped_row_offsets (TensorValue) – The grouped sequence offsets that use LoRA
max_lora_seq_len (int) – The maximum sequence length of any given LoRA in the batch
bias (TensorValue | None) – The LoRA bias
lora_end_idx (TensorValue)

Raises:

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

Return type:

`sgmv_lora_qkv_shrink()`

max.nn.kernels.sgmv_lora_qkv_shrink(input, lora_a, lora_ids, lora_grouped_offsets, lora_end_idx, max_lora_seq_len, max_rank)

source

LoRA shrink grouped matmul with planar Q/K/V output.

Performs the LoRA ‘shrink’ operation for routed tokens using SGMV (segmented grouped matrix-vector multiplication). Computes [M, K] @ [G, 3*rank, K]^T per active LoRA adapter, then permutes the flat [M, 3*rank] result into a planar layout [3, M, rank] representing separate Q, K, V projections.

Parameters:

input (TensorValue) – Routed activation matrix with shape (M, K), where M is the total number of tokens and K is the hidden dimension.
lora_a (TensorValue) – Shrink weights for all LoRA adapters, shape (G, 3*rank, K) where G is the number of adapters and rank is the LoRA rank.
lora_ids (TensorValue) – Expert/adapter indices for each active group, shape (num_active,). Values in range [0, G). May use -1 to indicate inactive slots.
lora_grouped_offsets (TensorValue) – Inclusive prefix sums of tokens per active adapter, shape (num_active + 1,). Defines per-adapter [start, end) ranges in input. Must be non-decreasing with offsets[0] == 0.
max_lora_seq_len (int) – Upper bound on tokens for any active adapter. Used for kernel tuning and memory allocation.
max_rank (int) – The maximum LoRA rank, determines output shape.
lora_end_idx (TensorValue)

Returns:

Output tensor with planar Q/K/V layout, shape (3, M, max_rank).

Raises:

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

Return type:

`sgmv_qkv_lora_fused()`

max.nn.kernels.sgmv_qkv_lora_fused(input, lora_a, lora_b, lora_ids, lora_grouped_offsets, lora_end_idx, q_dim, kv_dim, max_lora_seq_len, max_rank)

source

Returns the fused [q|k|v] LoRA contribution for a QKV projection.

Same shrink + single boundary-aware expand as sgmv_qkv_lora_kernel(), but returns the full [q|k|v] output of shape [M, q_dim + 2*kv_dim] (for the M LoRA tokens) instead of writing K/V into the cache. The caller adds it to the qkv projection before the fused rope + KV store.

Parameters:

input (TensorValue) – The input tensor.
lora_a (TensorValue) – The fused LoRA A tensor, shape (G, 3*R, K) (q/k/v stacked on the rank dim).
lora_b (TensorValue) – The fused LoRA B tensor, shape (G, q_dim + 2*kv_dim, R), with output rows concatenated along the projection dimension as Q | K | V.
lora_ids (TensorValue) – IDs of the LoRAs used for each sequence.
lora_grouped_offsets (TensorValue) – Grouped offsets for LoRA sequences.
lora_end_idx (TensorValue) – End index of LoRA tokens in the batch.
q_dim (int) – Output dimension of the Q projection.
kv_dim (int) – Output dimension of the K (and V) projection.
max_lora_seq_len (int) – The maximum sequence length of any LoRA in the batch.
max_rank (int) – The maximum rank for the LoRAs.

Raises:

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

Return type:

`sgmv_qkv_lora_kernel()`

max.nn.kernels.sgmv_qkv_lora_kernel(input, lora_a, lora_b, lora_ids, input_row_offsets, lora_grouped_offsets, lora_end_idx, batch_seq_len, kv_collection, kv_params, layer_idx, q_dim, kv_dim, max_lora_seq_len, max_rank)

source

Computes the SGMV QKV LoRA kernel for Q, K, V projections with LoRA.

The LoRA ‘expand’ (up-projection) is a single boundary-aware grouped matmul over one fused LoRA-B weight, correct under grouped-query attention (q_dim != kv_dim). The K/V projections are written directly into the paged KV cache; the Q projection is returned.

Parameters:

input (TensorValue) – The input tensor.
lora_a (TensorValue) – The fused LoRA A tensor, shape (G, 3*R, K) (q/k/v stacked on the rank dim).
lora_b (TensorValue) – The fused LoRA B tensor, shape (G, q_dim + 2*kv_dim, R), with output rows concatenated along the projection dimension as Q | K | V.
lora_ids (TensorValue) – IDs of the LoRAs used for each sequence.
input_row_offsets (TensorValue) – The sequence offsets that use LoRA.
lora_grouped_offsets (TensorValue) – Grouped offsets for LoRA sequences.
lora_end_idx (TensorValue) – End index of LoRA tokens in the batch.
batch_seq_len (TensorValue) – Total sequence length of the batch.
kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The KV cache.
kv_params (KVCacheParams) – The key-value cache configuration parameters.
layer_idx (TensorValue) – The layer index to retrieve the KV cache.
q_dim (int) – Output dimension of the Q projection.
kv_dim (int) – Output dimension of the K (and V) projection.
max_lora_seq_len (int) – The maximum sequence length of any given LoRA in the batch.
max_rank (int) – The maximum rank for the LoRAs.

Raises:

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

Return type:

`sleep()`

max.nn.kernels.sleep(duration_sec, device_ref)

source

Sleep for the given duration in seconds.

This kernel is supported on CPUs and GPUs. However, the timing may be completely inaccurate on AMD GPUs due to limitation of current time.sleep(…) impl.

Parameters:

duration_sec (BufferValue) – The duration to sleep in seconds.
device_ref (DeviceRef)

Return type:

None

`sliced_add()`

max.nn.kernels.sliced_add(x, y, lora_end_idx)

source

Adds tensors x and y element-wise for rows < lora_end_idx, otherwise copies x.

This is used for LoRA where only some sequences have LoRA applied. For rows in [0, lora_end_idx): c = x + y For rows in [lora_end_idx, batch_seq_len): c = x

Parameters:

x (TensorValue) – First input tensor.
y (TensorValue) – Second input tensor.
lora_end_idx (TensorValue) – End index of LoRA token portion (rows to apply add).

Return type:

`smallm_streaming_matmul()`

max.nn.kernels.smallm_streaming_matmul(a, b)

source

Computes bf16 a @ b^T with the MI355X small-M streaming kernel.

The kernel reads b in a fragment-major layout that this op permutes in the graph with preshuffle_block_scaled_b(); for a weight b the graph compiler runs that once, at model init. The streaming kernel serves M <= 32 (its measured win band) and larger batches fall back to generic matmul dispatch over the row-major b at execute time, so the op is never worse than the dispatch it replaces. Callers emit this op explicitly for decode-band vocab projections; nothing routes here through generic matmul dispatch.

Parameters:

a (TensorValue)
b (TensorValue)

Return type:

`spatial_merge()`

max.nn.kernels.spatial_merge(input, grid_thw, hidden_size, merge_size, *, out_rows)

source

Performs spatial merge operation on ragged input tensors.

This operation merges spatial dimensions of input patches according to the grid dimensions specified in grid_thw, replicating each grid item’s spatial rows across its t temporal slices.

Input and output row counts differ whenever any t > 1: the kernel consumes sum(h * w) rows and emits sum(t * h * w), so out_rows cannot be derived from input.shape. They coincide only at t == 1, which is why the two are separate arguments.

Parameters:

input (TensorValue) – Input tensor of shape [sum(h * w), hidden_size]
grid_thw (TensorValue) – Grid dimensions tensor of shape [batch_size, 3] containing [t, h, w] for each batch item, where:
- t: temporal/frame dimension
- h: height dimension
- w: width dimension
hidden_size (int) – Hidden dimension size
merge_size (int) – Size of spatial merge blocks (typically 2)
out_rows (int | str | Dim | integer | TypedAttr) – Row count of the output, sum(t * h * w). Usually the patch-row dimension of the caller’s pixel_values.

Returns:

Output tensor of shape [out_rows, hidden_size]

Raises:

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

Return type:

`store_k_cache_padded()`

max.nn.kernels.store_k_cache_padded(kv_collection, x_k, valid_lengths, layer_idx)

source

Stores the key tensor into the paged KV cache for padded inputs.

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection to write into.
x_k (TensorValue) – The key tensor of rank 4 containing the new key projections.
valid_lengths (TensorValue) – Buffer of shape [batch] (dtype uint32) indicating the actual (non-padded) sequence length for each batch element.
layer_idx (TensorValue) – The scalar layer index (dtype uint32) identifying which transformer layer’s cache to update.

Return type:

None

`store_k_cache_ragged()`

max.nn.kernels.store_k_cache_ragged(kv_collection, x_k, input_row_offsets, layer_idx)

source

Stores the key tensor into the paged KV cache for ragged inputs.

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection to write into.
x_k (TensorValue) – The key tensor of rank 3 containing the new key projections.
input_row_offsets (TensorValue) – Ragged tensor row offsets of shape [batch + 1] indicating where each sequence starts and ends. Must have dtype uint32.
layer_idx (TensorValue) – The scalar layer index (dtype uint32) identifying which transformer layer’s cache to update.

Return type:

None

`store_k_scale_cache_ragged()`

max.nn.kernels.store_k_scale_cache_ragged(kv_collection, x_k_scale, input_row_offsets, layer_idx, quantization_granularity)

source

Store key scale tensor into the paged KV cache.

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue])
x_k_scale (TensorValue)
input_row_offsets (TensorValue)
layer_idx (TensorValue)
quantization_granularity (int)

Return type:

None

`store_v_cache_ragged()`

max.nn.kernels.store_v_cache_ragged(kv_collection, x_v, input_row_offsets, layer_idx)

source

Stores the value tensor into the paged KV cache for ragged inputs.

Parameters:

kv_collection (KVCacheInputsPerDevice[TensorValue, BufferValue]) – The paged KV cache collection to write into.
x_v (TensorValue) – The value tensor of rank 3 containing the new value projections.
input_row_offsets (TensorValue) – Ragged tensor row offsets of shape [batch + 1] indicating where each sequence starts and ends. Must have dtype uint32.
layer_idx (TensorValue) – The scalar layer index (dtype uint32) identifying which transformer layer’s cache to update.

Return type:

None

`swiglu_quantize_dynamic_block_scaled()`

max.nn.kernels.swiglu_quantize_dynamic_block_scaled(gate_up, activation=None, tensor_sf=1.0, sf_vector_size=16, scales_type=float8_e4m3fn, out_type=uint8)

source

Quantizes the gated activation of a [gate | up] tensor with block scales.

One kernel computes what applying activation to each (gate, up) pair, rounding to bfloat16, and then quantize_dynamic_block_scaled() would, without writing the activation out. gate is the first half of each row and up the second. The formats and arguments match quantize_dynamic_block_scaled(). Runs on NVIDIA SM100 and SM120.

Parameters:

gate_up (TensorValue) – The bfloat16 input. Shape: [M, 2 * N].
activation (GatedActivation | None) – The gated activation, from GatedActivation’s fused_activation, alpha and limit. Defaults to plain SiLU.
tensor_sf (TensorValue | float) – The tensor-wise scale factor (inverted for NVFP4; identity for MXFP8).
sf_vector_size (int) – The block size for the scaling factors. 16 for NVFP4, 32 for MXFP4/MXFP8.
scales_type (DType) – The type of the scales tensor. float8_e4m3fn for NVFP4, float8_e8m0fnu for MXFP4/MXFP8.
out_type (DType) – The type of the output tensor. uint8 for packed FP4 or float8_e4m3fn for MXFP8.

Returns:

The quantized [M, N] activation ([M, N // 2] packed FP4) and its rank-5 interleaved scales.

Raises:

ValueError – If the GPU is not SM100 or SM120, gate_up is not a rank-2 bfloat16 tensor with a static even width whose halves are whole scale blocks, or the format arguments do not name a supported format.

Return type:

`top_k_per_row()`

max.nn.kernels.top_k_per_row(input, k, max_k)

source

top_k() over the last axis, with a count per row.

Row r gets the first k[r] picks ops.top_k(input, max_k) makes for it, in descending order; the rest of its max_k slots hold the lowest value of the dtype (-inf for floats) and index -1. Equal values may come out in any order, and which of them fill the last slots is unspecified. input must not contain NaN; mask it to -inf first if it may.

Parameters:

input (TensorValue) – [rows, n] values to select from.
k (TensorValue) – [rows] int64 picks per row on input’s device, each in [0, max_k].
max_k (int) – The output width, at most n.

Returns:

[rows, max_k] values in input’s dtype and their int64 column indices.

Return type:

`topk_fused_sampling()`

max.nn.kernels.topk_fused_sampling(logits, top_k, *, temperature=1.0, max_k=None, min_top_p=None, top_p=1.0, min_p=None, seed=0)

source

Performs top-k sampling with temperature scaling.

Parameters:

logits (TensorValue) – Input logits tensor of shape [batch_size, vocab_size].
top_k (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Number of top tokens to consider for sampling. Can be a scalar (which will be expanded to batch_size) or a tensor of shape [batch_size].
temperature (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Temperature for scaling logits before sampling.
max_k (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – Maximum value of k across the batch, or -1 if any row’s top_k is -1 or 0 (no limit): the kernel resolves a row’s -1 to max_k. Required when top_k is a tensor.
top_p (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Top-p (nucleus) sampling threshold. Can be a scalar or tensor.
min_p (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – Per-row min_p probability filtering threshold of shape [batch_size]. Tokens with probability below min_p * max_prob are zeroed before sampling.
seed (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Seed for the random number generator. Can be a scalar or tensor.
min_top_p (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None)

Returns:

Sampled tokens tensor of shape [batch_size, 1].

Raises:

ValueError – If input validation fails.

Return type:

`topk_fused_sampling_with_dist()`

max.nn.kernels.topk_fused_sampling_with_dist(logits, *, top_k, temperature, top_p, seed)

source

Samples a token per row and returns the distribution it came from.

Applies per-row temperature, top-k and top-p to logits, samples one token per row, and also returns the masked, renormalized distribution the token was drawn from. Speculative decoding subtracts that distribution to build its rejection residual, and reads the sampled token’s own probability out of it – a value the sampler must agree with, so it comes from the sampling kernel rather than a separate softmax.

This is the same kernel and the same code path as topk_fused_sampling(), with the distribution write enabled. GPU-only.

Parameters:

logits (TensorValue) – Raw logits [rows, vocab_size]. Softmax is fused in.
top_k (TensorValue) – Per-row top-k [rows], int64. -1 disables top-k, which is the production sentinel (SamplingParams normalizes 0 to -1).
temperature (TensorValue) – Per-row temperature [rows], float32. 0 is clamped, collapsing the row to its argmax.
top_p (TensorValue) – Per-row nucleus threshold [rows], float32.
seed (TensorValue) – Per-row RNG seed [rows], uint64.

Returns:

int64 [rows] and float32 [rows, vocab_size].

Return type:

(token_ids, distribution)

Raises:

ValueError – If the logits are not rank-2 on GPU, or a per-row parameter does not carry exactly one entry per logits row.

`topk_topp_masked_probs()`

max.nn.kernels.topk_topp_masked_probs(logits, *, top_k, temperature, top_p)

source

Computes each row’s top-k/top-p masked softmax, without sampling.

A token survives the joint constraint iff e_i = exp((logit_i - row_max) / temperature) exceeds the row’s cutoff, with masked probability e_i / kept_mass; everything else is zero. Each output row is that masked renormalized distribution – the same form topk_fused_sampling_with_dist() emits for the draft, so speculative verification reads target probabilities and builds its rejection residual from this one tensor with no in-graph rebuild.

GPU-only, with the same top-k sentinel and tie handling as topk_fused_sampling_with_dist().

Parameters:

logits (TensorValue) – Raw logits [rows, vocab_size].
top_k (TensorValue) – Per-row top-k [rows], int64.
temperature (TensorValue) – Per-row temperature [rows], float32.
top_p (TensorValue) – Per-row nucleus threshold [rows], float32.

Returns:

The masked distribution, float32 [rows, vocab_size].

Raises:

ValueError – If the logits are not rank-2 on GPU.

Return type:

`tpool_patch_merger()`

max.nn.kernels.tpool_patch_merger(input, grid_thws, kH, kW, max_h, max_w)

source

Performs temporal pooling patch merger on ragged video tokens.

For each video in the batch, averages the input across the temporal (T) dimension and rearranges the result according to the spatial merge kernel (kH, kW). Each video’s T*H*W input tokens are reduced to H*W output tokens. All videos are concatenated contiguously in the output.

Parameters:

input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Input tensor of shape [total_input_tokens, D] where total_input_tokens = sum(T_i * H_i * W_i) over all videos.
grid_thws (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Grid dimensions tensor of shape [n_videos, 3] with (T, H, W) per video. Must have dtype int64.
kH (int) – Merge kernel height.
kW (int) – Merge kernel width.
max_h (int | Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | float | integer[Any] | floating[Any] | DLPackArray) – Maximum H across all videos in the batch (for grid sizing). May be a Python int (baked as a graph constant) or a TensorValue computed at runtime (such as via ops.max).
max_w (int | Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | float | integer[Any] | floating[Any] | DLPackArray) – Maximum W across all videos in the batch (for grid sizing). May be a Python int or a TensorValue.

Returns:

Output tensor of shape [sum(H_i * W_i), D].

Raises:

ValueError – On invalid input shapes or dtypes.

Return type:

`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)

source

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 (KVCacheInputsPerDevice[TensorValue, BufferValue])
layer_idx (TensorValue)

Return type: