昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

SetBufferSize

功能说明

设置执行算子时默认分配显存的大小。仅配置默认大小，若不满足计算所需会在算子执行阶段分配更大空间。未配置时，默认是100M。不建议调整此默认值，分配过大显存会影响其它计算任务。

接口原型

torch_atb.set_buffer_size(bytes: int) -> None

参数

默认显存大小，单位为Byte，需要大于等于0。

返回值

无

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
import torch
import torch_atb

torch_atb.set_buffer_size(300000000)

# 后续内容为单算子调用示例
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])