SetBufferSize

Description

Sets the default size of the graphics memory allocated for operator execution. If the size does not suffice, a larger space will be allocated in the operator execution stage. If not otherwise specified, the default size is 100 MB. We recommend retaining the default value. Excess graphics memory may hinder other computation tasks.

API Prototype

torch_atb.set_buffer_size(bytes: int) -> None

Parameter

Default size of the graphics memory, in bytes. The value must be greater than or equal to 0.

Return Value

None

Example

import torch
import torch_atb

torch_atb.set_buffer_size(300000000)

# The following is a single-operator call example.
self_attention_param = torch_atb.SelfAttentionParam()
self_attention_param.head_num = 24
self_attention_param.kv_head_num = 24
self_attention_param.calc_type = torch_atb.SelfAttentionParam.CalcType.PA_ENCODER
self_attention = torch_atb.Operation(self_attention_param)

q = torch.ones(4096, 24, 64, dtype=torch.float16).npu()
k = torch.ones(4096, 24, 64, dtype=torch.float16).npu()
v = torch.ones(4096, 24, 64, dtype=torch.float16).npu()
seqlen = torch.tensor([4096], dtype=torch.int32)
intensors = [q,k,v,seqlen]

outputs = self_attention.forward([q,k,v,seqlen])

Parent topic: ATB Python API Usage Description