MatmulOptions#

class nvmath.linalg.advanced.MatmulOptions(
compute_type: int | None = None,
scale_type: int | None = None,
sm_count_target: int | None = 0,
fast_accumulation: bool | None = False,
device_id: int | None = None,
handle: int | None = None,
logger: Logger | None = None,
memory_limit: int | str | None = '80%',
blocking: Literal[True, 'auto'] = 'auto',
allocator: BaseCUDAMemoryManager | None = None,
)[source]#

用于为 Matmul 对象和包装函数 matmul() 提供选项的数据类。

compute_type#

CUDA 计算类型。一个合适的计算类型

类型:

nvmath.linalg.ComputeType

will be selected if not specified.
scale_type#

CUDA 数据类型。如果未指定，将选择与计算类型相符的合适数据类型。

类型:

nvmath.CudaDataType

sm_count_target#

用于执行的 SM 数量。默认值为 0，对应于所有可用的 SM。

类型:

int

fast_accumulation#

启用或禁用 FP8 快速累加模式。默认值为 False（禁用）。

类型:

bool

device_id#

CUDA 设备序号（如果 MM 操作数位于 CPU 上则使用）。如果未指定，将使用设备 0。

类型:

int | None

handle#

线性代数库句柄。如果未提供句柄，将创建一个。

类型:

int | None

logger#

Python Logger 对象。如果未提供 Logger 对象，将使用根 Logger。

类型:

logging.Logger

memory_limit#

MM 操作可用的最大内存。可以指定为值（带有可选后缀，如 K[iB]、M[iB]、G[iB]）或百分比。默认值为设备内存的 80%。

类型:

int | str | None

blocking#

一个标志，用于指定执行函数和方法的行为，例如 matmul() 和 Matmul.execute()。当 blocking 为 True 时，执行方法在操作完成之前不会返回。当 blocking 为 "auto" 时，如果输入在 GPU 上，则方法立即返回。当操作数在 CPU 上时，执行方法始终会阻塞，以确保用户不会在结果可用之前意外地使用它。默认值为 "auto"。

类型:

Literal[True, ‘auto’]

allocator#

一个支持 BaseCUDAMemoryManager 协议的对象，用于分配设备内存。如果未提供分配器，将使用来自库包的内存分配器（对于 PyTorch 操作数，使用 torch.cuda.caching_allocator_alloc()，否则使用 cupy.cuda.alloc()）。

类型:

nvmath.memory.BaseCUDAMemoryManager | None

另请参阅

Matmul, matmul()