SetBufferSize

Description

Sets the default size of the graphics memory allocated for operator execution. Only the default size is configured. If the size does not meet the computing requirements, larger space is allocated in the operator execution phase. If it is not configured, the default size is 100 MB. You are not advised to change the default value. If the allocated graphics memory is excessively large, other computing tasks will be affected.

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.

Returns

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