For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

mla_decode_sm100_dispatch

def mla_decode_sm100_dispatch[q_type: DType, k_t: MHAOperand, output_type: DType, mask_t: MHAMask, config: MHAConfig[config.dtype], depth: Int, num_heads: Int, group: Int = Int(1), *, ragged: Bool = False, _is_cache_length_accurate: Bool = False, decoding_warp_split_k: Bool = False, per_token_scale_rope_aware: Bool = False, sparse: Bool = False, rope_aware_kv_sparse: Bool = False, fold_shared_index: Bool = False](q: TileTensor[q_type, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: k_t, output: TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type], scale: Float32, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], mask: mask_t, scalar_args_buf: TileTensor[DType.int64, Storage=scalar_args_buf.Storage, linear_idx_type=scalar_args_buf.linear_idx_type], batch_size: Int, q_max_seq_len: Int, max_cache_valid_length: Int, ctx: DeviceContext, q_scale_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, indices_stride: Int = Int(0), topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, extra_k: OptionalReg[k_t] = None, extra_d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, extra_indices_stride: Int = Int(0), extra_topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, extra_scales_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, num_partitions_in: Optional[Int] = None)